Implement function for getting fields from inheritance chain

Author: theosotr

src/ir/ast.py

@@ -1,5 +1,5 @@
 # pylint: disable=dangerous-default-value
-from typing import List
+from typing import List, Set, Union
 from copy import deepcopy
 
 import src.ir.type_utils as tu
@@ -630,7 +630,7 @@ def get_overridable_fields(self):
             if f.can_override
         ]
 
-    def get_callable_functions(self, class_decls) -> List[FunctionDeclaration]:
+    def get_callable_functions(self, class_decls) -> Set[FunctionDeclaration]:
         """All functions that can be called in instantiations of this class
         """
         # Get functions that are implemented in the current class
@@ -668,6 +668,48 @@ def get_callable_functions(self, class_decls) -> List[FunctionDeclaration]:
 
         return functions
 
+    def get_all_fields(self, class_decls) -> Set[FieldDeclaration]:
+        """
+        All fields (including the inheritted ones) that can be accessed by
+        instantiations of this class.
+        """
+        fields = set(self.fields)
+        field_names = {f.name for f in fields}
+
+        if not self.superclasses:
+            return fields
+
+        # Retrieve fields from the inheritance chain.
+        super_cls = self.superclasses[0]
+        class_decl = tu.get_superclass_decl(super_cls, class_decls)
+
+        if not class_decl:
+            return fields
+
+        type_var_map = tu.get_superclass_type_var_map(super_cls, class_decl)
+
+        parent_fields = class_decl.get_all_fields(class_decls)
+
+        # substitute type variables in parent's functions
+        for f in parent_fields:
+            if f.name in field_names:
+                # We override this field in the current class
+                continue
+            new_f = deepcopy(f)
+            new_f.field_type = types.substitute_type(f.get_type(),
+                                                     type_var_map)
+            fields.add(new_f)
+
+        return fields
+
+    def get_all_attributes(self, class_decls) -> Set[Union[FunctionDeclaration, FieldDeclaration]]:
+        """
+        Get all attributes (fields + functions) from the inheritance chain
+        """
+        attributes = self.get_callable_functions(class_decls)
+        attributes.update(self.get_all_fields(class_decls))
+        return attributes
+
     def get_abstract_functions(self, class_decls) -> List[FunctionDeclaration]:
         # Get the abstract functions that are declared in the current class.
         functions = {

tests/test_ir.py

@@ -7,7 +7,7 @@
 from src.ir.kotlin_types import *
 
 
-def assert_abstract_funcs(actual, expected):
+def assert_declarations(actual, expected):
     assert len(actual) == len(expected)
     table = {f.name: f for f in expected}
     for f in actual:
@@ -183,8 +183,8 @@ def test_get_abstract_functions():
                             functions=[func2])
 
     assert cls1.get_abstract_functions([cls1, cls2]) == {func1}
-    assert_abstract_funcs(cls2.get_abstract_functions([cls1, cls2]),
-                          {func1, func2})
+    assert_declarations(cls2.get_abstract_functions([cls1, cls2]),
+                        {func1, func2})
 
 
 def test_get_abstract_functions_chain():
@@ -199,18 +199,17 @@ def test_get_abstract_functions_chain():
                             functions=[func2])
 
     assert cls1.get_abstract_functions([cls1, cls2, cls3]) == {func1}
-    assert_abstract_funcs(cls2.get_abstract_functions([cls1, cls2, cls3]),
-                          {func1})
-    assert_abstract_funcs(cls3.get_abstract_functions([cls1, cls2, cls3]),
-                          {func1, func2})
+    assert_declarations(cls2.get_abstract_functions([cls1, cls2, cls3]),
+                        {func1})
+    assert_declarations(cls3.get_abstract_functions([cls1, cls2, cls3]),
+                        {func1, func2})
 
     override_func = deepcopy(func1)
     func1.body = IntegerConstant(1, Integer)
     cls2.functions = [func1]
-    assert_abstract_funcs(cls2.get_abstract_functions([cls1, cls2, cls3]),
-                          [])
-    assert_abstract_funcs(cls3.get_abstract_functions([cls1, cls2, cls3]),
-                          {func2})
+    assert_declarations(cls2.get_abstract_functions([cls1, cls2, cls3]), [])
+    assert_declarations(cls3.get_abstract_functions([cls1, cls2, cls3]),
+                        {func2})
 
 
 def test_get_abstract_functions_parameterized():
@@ -232,8 +231,8 @@ def test_get_abstract_functions_parameterized():
     exp_func1.params[0].param_type = String
     exp_func1.ret_type = String
     exp_func1.inferred_type = String
-    assert_abstract_funcs(cls2.get_abstract_functions([cls1, cls2]),
-                          [exp_func1])
+    assert_declarations(cls2.get_abstract_functions([cls1, cls2]),
+                        [exp_func1])
 
 
 def test_get_abstract_functions_parameterized_chain():
@@ -262,5 +261,79 @@ def test_get_abstract_functions_parameterized_chain():
     exp_func1.params[0].param_type = actual_t
     exp_func1.ret_type = actual_t
     exp_func1.inferred_type = actual_t
-    assert_abstract_funcs(cls2.get_abstract_functions([cls1, cls2, cls3]),
-                          [exp_func1])
+    assert_declarations(cls2.get_abstract_functions([cls1, cls2, cls3]),
+                        [exp_func1])
+
+
+def test_get_fields():
+    field1 = FieldDeclaration("foo", String)
+    field2 = FieldDeclaration("bar", Any)
+    cls1 = ClassDeclaration("A", [], 0, fields=[field1])
+    cls2 = ClassDeclaration("B",
+                            [SuperClassInstantiation(cls1.get_type(), [])],
+                            fields=[field2])
+
+    assert cls1.get_all_fields([cls1, cls2]) == {field1}
+    assert_declarations(cls2.get_all_fields([cls1, cls2]),
+                        {field1, field2})
+
+    cls1 = ClassDeclaration("A", [], fields=[field1])
+    cls2 = ClassDeclaration("B",
+                            [SuperClassInstantiation(cls1.get_type(), [])],
+                            fields=[])
+    cls3 = ClassDeclaration("C",
+                            [SuperClassInstantiation(cls2.get_type(), [])],
+                            fields=[field2])
+
+    assert cls1.get_all_fields([cls1, cls2, cls3]) == {field1}
+    assert_declarations(cls2.get_all_fields([cls1, cls2, cls3]),
+                        {field1})
+    assert_declarations(cls3.get_all_fields([cls1, cls2, cls3]),
+                        {field1, field2})
+
+
+def test_get_fields_parameterized():
+    type_param1 = TypeParameter("T")
+    field1 = FieldDeclaration("foo", type_param1)
+    field2 = FieldDeclaration("bar", type_param1)
+    cls1 = ClassDeclaration("A", [],
+                            type_parameters=[type_param1],
+                            fields=[field1, field2])
+    cls2 = ClassDeclaration(
+        "B", [SuperClassInstantiation(cls1.get_type().new([String]), [])],
+        fields=[]
+    )
+
+    assert cls1.get_all_fields([cls1, cls2]) == {field1, field2}
+    exp_field1 = deepcopy(field1)
+    exp_field2 = deepcopy(field2)
+    exp_field1.field_type = String
+    exp_field2.field_type = String
+    assert_declarations(cls2.get_all_fields([cls1, cls2]),
+                        [field1, field2])
+
+
+    cls1 = ClassDeclaration("A", [],
+                            type_parameters=[type_param1],
+                            fields=[field1])
+
+    t_con = TypeConstructor("Foo", [TypeParameter("T")])
+    type_param2 = TypeParameter("T")
+    field2 = FieldDeclaration("bar", type_param2)
+    t = t_con.new([type_param2])
+    cls2 = ClassDeclaration(
+        "B", [SuperClassInstantiation(cls1.get_type().new([t]), [])],
+        type_parameters=[type_param2],
+        fields=[field2]
+    )
+    field3 = FieldDeclaration("bar", String)
+    cls3 = ClassDeclaration(
+        "C", [SuperClassInstantiation(cls2.get_type().new([String]), [])],
+        fields=[field3]
+    )
+    actual_t = t_con.new([String])
+    exp_field = deepcopy(field1)
+    exp_field.field_type = actual_t
+    assert_declarations(cls2.get_all_fields([cls1, cls2, cls3]),
+                        [exp_field, field3])
+