From b2def8aefd9daf9d120c47e6e4360752050d626e Mon Sep 17 00:00:00 2001
From: Troy Sornson <troy@sornson.io>
Date: Wed, 11 Dec 2024 21:38:53 -0700
Subject: [PATCH] Support types from default args that are constants

---
 spec/compiler/apply_types_spec.cr     | 19 +++++++++++++++++++
 src/compiler/crystal/command/typer.cr | 10 ++++++++--
 src/compiler/crystal/tools/typer.cr   | 21 ++++++++++++++-------
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/spec/compiler/apply_types_spec.cr b/spec/compiler/apply_types_spec.cr
index 155015e93b69..8a9d63d71bcd 100644
--- a/spec/compiler/apply_types_spec.cr
+++ b/spec/compiler/apply_types_spec.cr
@@ -425,6 +425,25 @@ describe Crystal::SourceTyper do
     OUTPUT
   end
 
+  it "types args with constant defaults" do
+    run_source_typer_spec(<<-INPUT, <<-OUTPUT, line_number: -1)
+    class Foo
+      MY_CONSTANT = 3
+      def test(arg = MY_CONSTANT); end
+    end
+
+    Foo.new.test(3.0)
+    INPUT
+    class Foo
+      MY_CONSTANT = 3
+
+      def test(arg : Float64 | Int32 = MY_CONSTANT) : Nil; end
+    end
+
+    Foo.new.test(3.0)
+    OUTPUT
+  end
+
   it "doesn't type methods that are inherited" do
     run_source_typer_spec(<<-INPUT, nil, line_number: -1)
     class Foo
diff --git a/src/compiler/crystal/command/typer.cr b/src/compiler/crystal/command/typer.cr
index 1c118ab605dd..0decce58c931 100644
--- a/src/compiler/crystal/command/typer.cr
+++ b/src/compiler/crystal/command/typer.cr
@@ -14,7 +14,7 @@ class Crystal::Command
     progress = false
     error_trace = false
 
-    OptionParser.parse(options) do |opts|
+    parser = OptionParser.new do |opts|
       opts.banner = <<-USAGE
         Usage: crystal tool apply-types [options] entrypoint [def_locator [def_locator [...]]]
 
@@ -69,7 +69,13 @@ class Crystal::Command
       end
     end
 
-    entrypoint = options.shift
+    parser.parse(options)
+
+    unless entrypoint = options.shift?
+      puts parser
+      exit
+    end
+
     def_locators = options
 
     results = SourceTyper.new(
diff --git a/src/compiler/crystal/tools/typer.cr b/src/compiler/crystal/tools/typer.cr
index 8be1ad67fa4a..861c85f185af 100644
--- a/src/compiler/crystal/tools/typer.cr
+++ b/src/compiler/crystal/tools/typer.cr
@@ -72,6 +72,10 @@ module Crystal
       # And now infer types of everything
       semantic_node = program.semantic nodes, cleanup: true
 
+      # We might run semantic later in an attempt to resolve defaults, don't display those stats or progress
+      @program.progress_tracker.stats = false
+      @program.progress_tracker.progress = false
+
       # Use the DefVisitor to locate and match any 'def's that match a def_locator
       def_visitor = DefVisitor.new(@def_locators, @excludes, entrypoint)
       semantic_node.accept(def_visitor)
@@ -264,13 +268,6 @@ module Crystal
             else
               raise "Unknown handling of arg #{arg} at #{def_instance.location} in #{def_instance}\n#{parsed}"
             end
-
-            # Special case - we can have default args that are never used be a different type than what was set.
-            # Ensure those default arg types also get respected (i.e. `arg = nil` => `arg : Int32? = nil` instead
-            # of `arg : Int32 = nil`)
-            if def_val = arg.default_value
-              all_typed_args[arg.external_name] << program.semantic(def_val).type
-            end
           end
 
           encountered_non_splat_arg_def_instance |= !encountered_splat_arg
@@ -281,6 +278,16 @@ module Crystal
           end
         end
 
+        parsed.args.each do |arg|
+          if def_val = arg.default_value
+            if def_val.to_s.matches?(/^[A-Z_]+$/)
+              # This looks like a constant, let's try qualifying it with the parent type
+              def_val = Crystal::Path.new([parsed.owner.to_s, def_val.to_s])
+            end
+            all_typed_args[arg.external_name] << program.semantic(def_val).type rescue nil
+          end
+        end
+
         # If a given collection of def_instances has a splat defined AND at least one def_instance didn't have a type for it,
         # then we can't add types to the signature.
         # https://crystal-lang.org/reference/1.14/syntax_and_semantics/type_restrictions.html#splat-type-restrictions