Add ArrayLiteral#[]= in macros

crystal-lang · Nov 30, 2016 · c5cf484 · c5cf484
1 parent 17c829d
commit c5cf484
Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/spec/compiler/macro/macro_methods_spec.cr b/spec/compiler/macro/macro_methods_spec.cr
@@ -634,6 +634,10 @@ describe "macro methods" do
       assert_macro "", %({{ [1, 2, 3, 4, 5][1, 3] }}), [] of ASTNode, %([2, 3, 4])
     end
 
+    it "executes []=" do
+      assert_macro "", %({% a = [0]; a[0] = 2 %}{{a[0]}}), [] of ASTNode, "2"
+    end
+
     it "executes of" do
       assert_macro "x", %({{ x.of }}), [ArrayLiteral.new([] of ASTNode, of: Path.new(["Int64"]))] of ASTNode, %(Int64)
     end

diff --git a/src/compiler/crystal/macros/methods.cr b/src/compiler/crystal/macros/methods.cr
@@ -1810,6 +1810,22 @@ private def intepret_array_or_tuple_method(object, klass, method, args, block, i
     else
       object.wrong_number_of_arguments "#{klass}#[]", args.size, 1
     end
+  when "[]="
+    object.interpret_two_args_method(method, args) do |index_node, value|
+      unless index_node.is_a?(Crystal::NumberLiteral)
+        index_node.raise "expected index argument to ArrayLiteral#[]= to be a number, not #{index_node.class_desc}"
+      end
+
+      index = index_node.to_number.to_i
+      index += object.elements.size if index < 0
+
+      unless 0 <= index < object.elements.size
+        index_node.raise "index out of bounds (index: #{index}, size: #{object.elements.size}"
+      end
+
+      object.elements[index] = value
+      value
+    end
   when "unshift"
     case args.size
     when 1