From 17f538d8cc2562995b9881a2f9fe7f75cd204afc Mon Sep 17 00:00:00 2001
From: "Johan B.W. de Vries" <jbwdevries@gmail.com>
Date: Sat, 24 Dec 2022 19:25:30 +0100
Subject: [PATCH] Fixed a missing constraint

---
 phasm/type3/constraints.py                    |  9 ++++++---
 phasm/type3/constraintsgenerator.py           |  2 +-
 .../test_lang/test_static_array.py            | 19 ++++++++++++++++++-
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/phasm/type3/constraints.py b/phasm/type3/constraints.py
index 0cde42c..06e0dfc 100644
--- a/phasm/type3/constraints.py
+++ b/phasm/type3/constraints.py
@@ -350,15 +350,17 @@ class CanBeSubscriptedConstraint(ConstraintBase):
     """
     A value that is subscipted, i.e. a[0] (tuple) or a[b] (static array)
     """
-    __slots__ = ('type3', 'index', 'index_type3', )
+    __slots__ = ('ret_type3', 'type3', 'index', 'index_type3', )
 
+    ret_type3: types.Type3OrPlaceholder
     type3: types.Type3OrPlaceholder
     index: ourlang.Expression
     index_type3: types.Type3OrPlaceholder
 
-    def __init__(self, type3: types.Type3OrPlaceholder, index: ourlang.Expression, comment: Optional[str] = None) -> None:
+    def __init__(self, ret_type3: types.Type3OrPlaceholder, type3: types.Type3OrPlaceholder, index: ourlang.Expression, comment: Optional[str] = None) -> None:
         super().__init__(comment=comment)
 
+        self.ret_type3 = ret_type3
         self.type3 = type3
         self.index = index
         self.index_type3 = index.type3
@@ -373,7 +375,8 @@ class CanBeSubscriptedConstraint(ConstraintBase):
         if isinstance(self.type3, types.AppliedType3):
             if self.type3.base == types.static_array:
                 return [
-                    SameTypeConstraint(types.u32, self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b')
+                    SameTypeConstraint(types.u32, self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
+                    SameTypeConstraint(self.type3.args[0], self.ret_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
                 ]
 
         # FIXME: bytes
diff --git a/phasm/type3/constraintsgenerator.py b/phasm/type3/constraintsgenerator.py
index 1ffc610..241d0ec 100644
--- a/phasm/type3/constraintsgenerator.py
+++ b/phasm/type3/constraintsgenerator.py
@@ -103,7 +103,7 @@ def expression(ctx: Context, inp: ourlang.Expression) -> Generator[ConstraintBas
         yield from expression(ctx, inp.varref)
         yield from expression(ctx, inp.index)
 
-        yield CanBeSubscriptedConstraint(inp.varref.type3, inp.index)
+        yield CanBeSubscriptedConstraint(inp.type3, inp.varref.type3, inp.index)
         return
 
     if isinstance(inp, ourlang.AccessStructMember):
diff --git a/tests/integration/test_lang/test_static_array.py b/tests/integration/test_lang/test_static_array.py
index 7b3ce41..0a43eed 100644
--- a/tests/integration/test_lang/test_static_array.py
+++ b/tests/integration/test_lang/test_static_array.py
@@ -96,7 +96,24 @@ def helper(array: {type_}[3]) -> {type_}:
     assert TYPE_MAP[type_] == type(result.returned_value)
 
 @pytest.mark.integration_test
-def test_function_call_element():
+def test_function_call_element_ok():
+    code_py = """
+CONSTANT: u64[3] = (250, 250000, 250000000, )
+
+@exported
+def testEntry() -> u64:
+    return helper(CONSTANT[0])
+
+def helper(x: u64) -> u64:
+    return x
+"""
+
+    result = Suite(code_py).run_code()
+
+    assert 250 == result.returned_value
+
+@pytest.mark.integration_test
+def test_function_call_element_type_mismatch():
     code_py = """
 CONSTANT: u64[3] = (250, 250000, 250000000, )