Added some RangeLiteral methods in macros: begin, end, excludes_end? …

…and map
crystal-lang · Jun 9, 2016 · 2e01a89 · 2e01a89
1 parent 8161824
commit 2e01a89
Showing 3 changed files with 75 additions and 0 deletions.
diff --git a/spec/compiler/macro/macro_methods_spec.cr b/spec/compiler/macro/macro_methods_spec.cr
@@ -977,6 +977,25 @@ describe "macro methods" do
     end
   end
 
+  describe "range methods" do
+    it "executes begin" do
+      assert_macro "x", %({{x.begin}}), [RangeLiteral.new(1.int32, 2.int32, true)] of ASTNode, "1"
+    end
+
+    it "executes end" do
+      assert_macro "x", %({{x.end}}), [RangeLiteral.new(1.int32, 2.int32, true)] of ASTNode, "2"
+    end
+
+    it "executes excludes_end?" do
+      assert_macro "x", %({{x.excludes_end?}}), [RangeLiteral.new(1.int32, 2.int32, true)] of ASTNode, "true"
+    end
+
+    it "executes map" do
+      assert_macro "x", %({{x.map(&.stringify)}}), [RangeLiteral.new(1.int32, 3.int32, false)] of ASTNode, %(["1", "2", "3"])
+      assert_macro "x", %({{x.map(&.stringify)}}), [RangeLiteral.new(1.int32, 3.int32, true)] of ASTNode, %(["1", "2"])
+    end
+  end
+
   describe "env" do
     it "has key" do
       ENV["FOO"] = "foo"

diff --git a/src/compiler/crystal/macros.cr b/src/compiler/crystal/macros.cr
@@ -555,6 +555,22 @@ module Crystal::Macros
 
   # A range literal.
   class RangeLiteral < ASTNode
+    # Similar to `Range#begin`
+    def begin : ASTNode
+    end
+
+    # Similar to `Range#end`
+    def end : ASTNode
+    end
+
+    # Similar to `Range#excludes_end?`
+    def excludes_end? : ASTNode
+    end
+
+    # Similar to `Enumerable#map` for a `Range`.
+    # Only works on ranges of `NumberLiteral`s considered as integers.
+    def map : ArrayLiteral
+    end
   end
 
   # A regex literal.

diff --git a/src/compiler/crystal/macros/methods.cr b/src/compiler/crystal/macros/methods.cr
@@ -706,6 +706,46 @@ module Crystal
     end
   end
 
+  class RangeLiteral
+    def interpret(method, args, block, interpreter)
+      case method
+      when "begin"
+        interpret_argless_method(method, args) { self.from }
+      when "end"
+        interpret_argless_method(method, args) { self.to }
+      when "excludes_end?"
+        interpret_argless_method(method, args) { BoolLiteral.new(self.exclusive) }
+      when "map"
+        interpret_argless_method(method, args) do
+          raise "map expects a block" unless block
+
+          from = self.from
+          to = self.to
+
+          unless from.is_a?(NumberLiteral)
+            raise "range begin must be a NumberLiteral, not #{from.class_desc}"
+          end
+
+          unless to.is_a?(NumberLiteral)
+            raise "range end must be a NumberLiteral, not #{to.class_desc}"
+          end
+
+          from = from.to_number.to_i64
+          to = to.to_number.to_i64
+          block_arg = block.args.first?
+
+          range = self.exclusive ? (from...to) : (from..to)
+          ArrayLiteral.map(range) do |num|
+            interpreter.define_var(block_arg.name, NumberLiteral.new(num)) if block_arg
+            interpreter.accept block.body
+          end
+        end
+      else
+        super
+      end
+    end
+  end
+
   class MetaVar < ASTNode
     def to_macro_id
       @name