Add reduce to ArrayLiteral for macros (#5294)

crystal-lang · Nov 18, 2017 · 818c438 · 818c438
1 parent 47f7e35
commit 818c438
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/spec/compiler/macro/macro_methods_spec.cr b/spec/compiler/macro/macro_methods_spec.cr
@@ -555,6 +555,14 @@ describe "macro methods" do
       assert_macro "", %({{[1, 2, 3].map { |e| e == 2 }}}), [] of ASTNode, "[false, true, false]"
     end
 
+    it "executes reduce with no initial value" do
+      assert_macro "", %({{[1, 2, 3].reduce { |acc, val| acc * val }}}), [] of ASTNode, "6"
+    end
+
+    it "executes reduce with initial value" do
+      assert_macro "", %({{[1, 2, 3].reduce(4) { |acc, val| acc * val }}}), [] of ASTNode, "24"
+    end
+
     it "executes map with constants" do
       assert_macro "x", %({{x.map { |e| e.id }}}), [ArrayLiteral.new([Path.new("Foo"), Path.new("Bar")] of ASTNode)] of ASTNode, "[Foo, Bar]"
     end

diff --git a/src/compiler/crystal/macros.cr b/src/compiler/crystal/macros.cr
@@ -582,6 +582,10 @@ module Crystal::Macros
     def reject(&block) : ArrayLiteral
     end
 
+    # Similar to `Enumerable#reduce`
+    def reduce(&block) : ASTNode
+    end
+
     # Similar to `Array#shuffle`
     def shuffle : ArrayLiteral
     end

diff --git a/src/compiler/crystal/macros/methods.cr b/src/compiler/crystal/macros/methods.cr
@@ -1963,6 +1963,30 @@ private def intepret_array_or_tuple_method(object, klass, method, args, block, i
       raise "reject expects a block" unless block
       filter(object, klass, block, interpreter, keep: false)
     end
+  when "reduce"
+    raise "reduce expects a block" unless block
+    accumulate_arg = block.args.first?
+    value_arg = block.args[1]?
+    case args.size
+    when 0
+      object.interpret_argless_method(method, args) do
+        object.elements.reduce do |accumulate, elem|
+          interpreter.define_var(accumulate_arg.name, accumulate) if accumulate_arg
+          interpreter.define_var(value_arg.name, elem) if value_arg
+          interpreter.accept block.body
+        end
+      end
+    when 1
+      object.interpret_one_arg_method(method, args) do |arg|
+        object.elements.reduce(arg) do |accumulate, elem|
+          interpreter.define_var(accumulate_arg.name, accumulate) if accumulate_arg
+          interpreter.define_var(value_arg.name, elem) if value_arg
+          interpreter.accept block.body
+        end
+      end
+    else
+      raise "only 0 or 1 args expected for reduce, got #{args.size}"
+    end
   when "shuffle"
     klass.new(object.elements.shuffle)
   when "sort"