Merge pull request #2133 from herwinw/for_ivar

Support more targets in for loops

lib/natalie/compiler/args.rb

@@ -46,6 +46,12 @@ def transform_arg(arg)
         when ::Prism::ArrayNode
           clean_up_keyword_args
           transform_destructured_arg(arg)
+        when ::Prism::InstanceVariableTargetNode
+          transform_instance_variable_arg(arg)
+        when ::Prism::ClassVariableTargetNode
+          transform_class_variable_arg(arg)
+        when ::Prism::ConstantTargetNode
+          transform_constant_arg(arg)
         when ::Prism::RequiredParameterNode
           clean_up_keyword_args
           transform_required_arg(arg)
@@ -158,6 +164,22 @@ def transform_destructured_arg(arg)
         @instructions << PopInstruction.new
       end
 
+      def transform_instance_variable_arg(arg)
+        @instructions << ArrayShiftInstruction.new
+        @instructions << InstanceVariableSetInstruction.new(arg.name)
+      end
+
+      def transform_class_variable_arg(arg)
+        @instructions << ArrayShiftInstruction.new
+        @instructions << ClassVariableSetInstruction.new(arg.name)
+      end
+
+      def transform_constant_arg(arg)
+        @instructions << ArrayShiftInstruction.new
+        @instructions << PushSelfInstruction.new
+        @instructions << ConstSetInstruction.new(arg.name)
+      end
+
       def transform_rest_arg(arg)
         if (name = arg.name)
           @instructions << variable_set(name)

lib/natalie/compiler/pass1.rb

@@ -2665,6 +2665,9 @@ def maximum_arg_count(args)
 
         args.count do |arg|
           %i[
+            class_variable_target_node
+            constant_target_node
+            instance_variable_target_node
             local_variable_target_node
             multi_target_node
             optional_parameter_node

spec/language/for_spec.rb

@@ -69,42 +69,38 @@ def obj.each
   it "allows an instance variable as an iterator name" do
     m = [1,2,3]
     n = 0
-    #for @var in m
-      #n += 1
-    #end
-    #@var.should == 3
-    NATFIXME 'Support Prism::InstanceVariableTargetNode', exception: SpecFailedException do
-      n.should == 3
+    for @var in m
+      n += 1
     end
+    @var.should == 3
+    n.should == 3
   end
 
   it "allows a class variable as an iterator name" do
     class OFor
       m = [1,2,3]
       n = 0
-      #for @@var in m
-        #n += 1
-      #end
-      #@@var.should == 3
-      NATFIXME 'Support Prism::ClassVariableTargetNode', exception: SpecFailedException do
-        n.should == 3
+      for @@var in m
+        n += 1
       end
+      @@var.should == 3
+      n.should == 3
     end
   end
 
   it "allows a constant as an iterator name" do
     class OFor
       m = [1,2,3]
       n = 0
-      NATFIXME 'Support Prism::ConstantTargetNode', exception: SpecFailedException do
+      NATFIXME 'Complain when assigning to initialized constant', exception: SpecFailedException, message: /should have printed a warning/ do
         -> {
-          #for CONST in m
-            #n += 1
-          #end
+          for CONST in m
+            n += 1
+          end
         }.should complain(/already initialized constant/)
-        CONST.should == 3
-        n.should == 3
       end
+      CONST.should == 3
+      n.should == 3
     end
   end