From bfa64d463acd35e0037e3ca9fd10eb2d3db82a05 Mon Sep 17 00:00:00 2001
From: Quinton Miller <nicetas.c@gmail.com>
Date: Wed, 14 Aug 2024 02:34:48 +0800
Subject: [PATCH 1/2] Implement `#sort_by` inside macros using
 `Enumerable#sort_by`

---
 spec/compiler/macro/macro_methods_spec.cr | 14 ++++++++++----
 src/compiler/crystal/macros/methods.cr    | 22 ++++++++++++++--------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/spec/compiler/macro/macro_methods_spec.cr b/spec/compiler/macro/macro_methods_spec.cr
index 38b08f44568a..385e165a3504 100644
--- a/spec/compiler/macro/macro_methods_spec.cr
+++ b/spec/compiler/macro/macro_methods_spec.cr
@@ -928,6 +928,16 @@ module Crystal
         assert_macro %({{["c".id, "b", "a".id].sort}}), %([a, "b", c])
       end
 
+      it "executes sort_by" do
+        assert_macro %({{["abc", "a", "ab"].sort_by { |x| x.size }}}), %(["a", "ab", "abc"])
+      end
+
+      it "calls block exactly once for each element in #sort_by" do
+        assert_macro <<-CRYSTAL, %(5)
+          {{ (i = 0; ["abc", "a", "ab", "abcde", "abcd"].sort_by { i += 1 }; i) }}
+          CRYSTAL
+      end
+
       it "executes uniq" do
         assert_macro %({{[1, 1, 1, 2, 3, 1, 2, 3, 4].uniq}}), %([1, 2, 3, 4])
       end
@@ -1020,10 +1030,6 @@ module Crystal
         assert_macro %({{{:a => 1, :b => 3}.size}}), "2"
       end
 
-      it "executes sort_by" do
-        assert_macro %({{["abc", "a", "ab"].sort_by { |x| x.size }}}), %(["a", "ab", "abc"])
-      end
-
       it "executes empty?" do
         assert_macro %({{{:a => 1}.empty?}}), "false"
       end
diff --git a/src/compiler/crystal/macros/methods.cr b/src/compiler/crystal/macros/methods.cr
index d3a1a1cc15a6..8cf557ee7b77 100644
--- a/src/compiler/crystal/macros/methods.cr
+++ b/src/compiler/crystal/macros/methods.cr
@@ -554,7 +554,8 @@ module Crystal
     end
 
     def interpret_compare(other : NumberLiteral)
-      to_number <=> other.to_number
+      # it should not be possible to obtain a NaN number literal
+      (to_number <=> other.to_number).not_nil!
     end
 
     def bool_bin_op(method, args, named_args, block, &)
@@ -3232,12 +3233,17 @@ end
 private def sort_by(object, klass, block, interpreter)
   block_arg = block.args.first?
 
-  klass.new(object.elements.sort { |x, y|
-    block_arg.try { |arg| interpreter.define_var(arg.name, x) }
-    x_result = interpreter.accept(block.body)
-    block_arg.try { |arg| interpreter.define_var(arg.name, y) }
-    y_result = interpreter.accept(block.body)
+  klass.new(object.elements.sort_by do |elem|
+    block_arg.try { |arg| interpreter.define_var(arg.name, elem) }
+    result = interpreter.accept(block.body)
+    InterpretCompareWrapper.new(result)
+  end)
+end
 
-    x_result.interpret_compare(y_result)
-  })
+private record InterpretCompareWrapper, node : Crystal::ASTNode do
+  include Comparable(self)
+
+  def <=>(other : self) : Int
+    node.interpret_compare(other.node)
+  end
 end

From 3bedc08a167d04349ef323b244553863728dcb99 Mon Sep 17 00:00:00 2001
From: Quinton Miller <nicetas.c@gmail.com>
Date: Wed, 14 Aug 2024 02:42:11 +0800
Subject: [PATCH 2/2] fixup

---
 src/compiler/crystal/macros/methods.cr | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/compiler/crystal/macros/methods.cr b/src/compiler/crystal/macros/methods.cr
index 8cf557ee7b77..8a7aa569fa95 100644
--- a/src/compiler/crystal/macros/methods.cr
+++ b/src/compiler/crystal/macros/methods.cr
@@ -554,8 +554,7 @@ module Crystal
     end
 
     def interpret_compare(other : NumberLiteral)
-      # it should not be possible to obtain a NaN number literal
-      (to_number <=> other.to_number).not_nil!
+      to_number <=> other.to_number
     end
 
     def bool_bin_op(method, args, named_args, block, &)
@@ -3243,7 +3242,7 @@ end
 private record InterpretCompareWrapper, node : Crystal::ASTNode do
   include Comparable(self)
 
-  def <=>(other : self) : Int
+  def <=>(other : self)
     node.interpret_compare(other.node)
   end
 end