Merge pull request #324 from OpShin/feat/index_list

Implement list index
OpShin · Jan 31, 2024 · 1d71d09 · 1d71d09
2 parents 584722f + 59bbb35
commit 1d71d09
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/opshin/tests/test_stdlib.py b/opshin/tests/test_stdlib.py
@@ -46,6 +46,28 @@ def validator(x: Dict[int, bytes], y: int) -> bytes:
             exp = None
         self.assertEqual(ret, exp, "dict[] returned wrong value")
 
+    @given(st.data())
+    def test_list_index(self, data):
+        source_code = """
+def validator(x: List[int], z: int) -> int:
+    return x.index(z)
+            """
+        xs = data.draw(st.lists(st.integers()))
+        z = data.draw(
+            st.one_of(st.sampled_from(xs), st.integers())
+            if len(xs) > 0
+            else st.integers()
+        )
+        try:
+            ret = eval_uplc_value(source_code, xs, z)
+        except RuntimeError as e:
+            ret = None
+        try:
+            exp = xs.index(z)
+        except ValueError:
+            exp = None
+        self.assertEqual(ret, exp, "list.index returned wrong value")
+
     @given(xs=st.dictionaries(st.integers(), st.binary()))
     def test_dict_keys(self, xs):
         source_code = """

diff --git a/opshin/types.py b/opshin/types.py
@@ -874,6 +874,48 @@ class ListType(ClassType):
     def __ge__(self, other):
         return isinstance(other, ListType) and self.typ >= other.typ
 
+    def attribute_type(self, attr) -> "Type":
+        if attr == "index":
+            return InstanceType(
+                FunctionType(frozenlist([self.typ]), IntegerInstanceType)
+            )
+        super().attribute_type(attr)
+
+    def attribute(self, attr) -> plt.AST:
+        if attr == "index":
+            return OLambda(
+                ["self", "x"],
+                OLet(
+                    [("x", plt.Force(OVar("x")))],
+                    plt.Apply(
+                        plt.RecFun(
+                            OLambda(
+                                ["index", "xs", "a"],
+                                plt.IteNullList(
+                                    OVar("xs"),
+                                    plt.TraceError("Did not find element in list"),
+                                    plt.Ite(
+                                        plt.EqualsInteger(
+                                            OVar("x"), plt.HeadList(OVar("xs"))
+                                        ),
+                                        OVar("a"),
+                                        plt.Apply(
+                                            OVar("index"),
+                                            OVar("index"),
+                                            plt.TailList(OVar("xs")),
+                                            plt.AddInteger(OVar("a"), plt.Integer(1)),
+                                        ),
+                                    ),
+                                ),
+                            ),
+                        ),
+                        OVar("self"),
+                        plt.Integer(0),
+                    ),
+                ),
+            )
+        super().attribute(attr)
+
     def stringify(self, recursive: bool = False) -> plt.AST:
         return OLambda(
             ["self"],