From aa6685a0cbb5757171017165ca5a6f4ca517e1bf Mon Sep 17 00:00:00 2001
From: Torsten Ruger <torsten@villataika.fi>
Date: Tue, 20 Dec 2016 15:46:23 +0200
Subject: [PATCH] test collector for ivar assignment

---
 lib/melon/type_collector.rb            | 10 ++++++++-
 test/melon/compiler/test_all.rb        |  1 +
 test/melon/compiler/test_class.rb      |  8 +------
 test/melon/compiler/test_clollector.rb | 29 ++++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 8 deletions(-)
 create mode 100644 test/melon/compiler/test_clollector.rb

diff --git a/lib/melon/type_collector.rb b/lib/melon/type_collector.rb
index 42e3e235..c5c3921f 100644
--- a/lib/melon/type_collector.rb
+++ b/lib/melon/type_collector.rb
@@ -12,10 +12,18 @@ module Melon
     end
 
     def on_ivar(statement)
+      add_ivar(statement)
+    end
+
+    def on_ivasgn( statement )
+      add_ivar(statement)
+    end
+
+    def add_ivar(statement)
       var = statement.children[0].to_s[1..-1].to_sym
       @ivars[var] = :Object #guess, can maybe guess better
     end
-
+    
     def handler_missing(node)
       node.children.each do |kid |
         process(kid) if kid.is_a?(AST::Node)
diff --git a/test/melon/compiler/test_all.rb b/test/melon/compiler/test_all.rb
index 9e6f195a..d62ab71f 100644
--- a/test/melon/compiler/test_all.rb
+++ b/test/melon/compiler/test_all.rb
@@ -1 +1,2 @@
 require_relative "test_class"
+require_relative "test_collector"
diff --git a/test/melon/compiler/test_class.rb b/test/melon/compiler/test_class.rb
index 44a1a804..30db5487 100644
--- a/test/melon/compiler/test_class.rb
+++ b/test/melon/compiler/test_class.rb
@@ -14,16 +14,10 @@ module Melon
     end
 
     def test_creates_class_with_deriviation
-     Compiler.compile  "class Test2 < List ;end"
+      Compiler.compile  "class Test2 < List ;end"
       assert t = Parfait::Space.object_space.get_class_by_name(:Test2) , "No classes created"
       assert_equal :List , t.super_class_name
     end
 
-    def test_picks_up_an_ivar_name
-      Compiler.compile  "class TestIvar < Object ; def meth; return @ivar;end; end"
-      itest = Parfait::Space.object_space.get_class_by_name(:TestIvar)
-      assert itest.instance_type.instance_names.include?(:ivar) , itest.instance_type.instance_names.inspect
-    end
-
   end
 end
diff --git a/test/melon/compiler/test_clollector.rb b/test/melon/compiler/test_clollector.rb
new file mode 100644
index 00000000..54beed95
--- /dev/null
+++ b/test/melon/compiler/test_clollector.rb
@@ -0,0 +1,29 @@
+require_relative "helper"
+
+module Melon
+  class TestCollector < MiniTest::Test
+
+    def setup
+      Register.machine.boot unless Register.machine.booted
+    end
+
+    def test_ivar_name
+      Compiler.compile  "class TestIvar < Object ; def meth; @ivar;end; end"
+      itest = Parfait::Space.object_space.get_class_by_name(:TestIvar)
+      assert itest.instance_type.instance_names.include?(:ivar) , itest.instance_type.instance_names.inspect
+    end
+
+    def test_ivar_assign
+      Compiler.compile  "class TestIvar2 < Object ; def meth; @ivar = 5 ;end; end"
+      itest = Parfait::Space.object_space.get_class_by_name(:TestIvar2)
+      assert itest.instance_type.instance_names.include?(:ivar) , itest.instance_type.instance_names.inspect
+    end
+
+    def test_ivar_operator_assign
+      Compiler.compile  "class TestIvar3 < Object ; def meth; @ivar += 5 ;end; end"
+      itest = Parfait::Space.object_space.get_class_by_name(:TestIvar3)
+      assert itest.instance_type.instance_names.include?(:ivar) , itest.instance_type.instance_names.inspect
+    end
+
+  end
+end