Merge pull request #397 from SCMusson/feat_unions

Unions of PlutusData and BuiltinData
OpShin · Aug 24, 2024 · e927d1d · e927d1d
2 parents cc234a1 + 22bec8f
commit e927d1d
Show file tree

Hide file tree

Showing 4 changed files with 479 additions and 13 deletions.
diff --git a/opshin/compiler.py b/opshin/compiler.py
@@ -422,7 +422,14 @@ def visit_Call(self, node: TypedCall) -> plt.AST:
             bind_self = node.func.typ.typ.bind_self
         bound_vs = sorted(list(node.func.typ.typ.bound_vars.keys()))
         args = []
-        for a, t in zip(node.args, node.func.typ.typ.argtyps):
+        for i, (a, t) in enumerate(zip(node.args, node.func.typ.typ.argtyps)):
+            # now impl_from_args has been chosen, skip type arg
+            if (
+                hasattr(node.func, "orig_id")
+                and node.func.orig_id == "isinstance"
+                and i == 1
+            ):
+                continue
             assert isinstance(t, InstanceType)
             # pass in all arguments evaluated with the statemonad
             a_int = self.visit(a)

diff --git a/opshin/fun_impls.py b/opshin/fun_impls.py
@@ -87,6 +87,80 @@ def impl_from_args(self, args: typing.List[Type]) -> plt.AST:
         return print
 
 
+class IsinstanceImpl(PolymorphicFunction):
+    def type_from_args(self, args: typing.List[Type]) -> FunctionType:
+        assert (
+            len(args) == 2
+        ), f"isinstance takes two arguments [object, type], but {len(args)} were given"
+        return FunctionType(args, BoolInstanceType)
+
+    def impl_from_args(self, args: typing.List[Type]) -> plt.AST:
+        if isinstance(args[1], IntegerType):
+            return OLambda(
+                ["x"],
+                plt.ChooseData(
+                    OVar("x"),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(True),
+                    plt.Bool(False),
+                ),
+            )
+        elif isinstance(args[1], ByteStringType):
+            return OLambda(
+                ["x"],
+                plt.ChooseData(
+                    OVar("x"),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(True),
+                ),
+            )
+        elif isinstance(args[1], RecordType):
+            return OLambda(
+                ["x"],
+                plt.ChooseData(
+                    OVar("x"),
+                    plt.Bool(True),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                ),
+            )
+        elif isinstance(args[1], ListType):
+            return OLambda(
+                ["x"],
+                plt.ChooseData(
+                    OVar("x"),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(True),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                ),
+            )
+        elif isinstance(args[1], DictType):
+            return OLambda(
+                ["x"],
+                plt.ChooseData(
+                    OVar("x"),
+                    plt.Bool(False),
+                    plt.Bool(True),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                    plt.Bool(False),
+                ),
+            )
+        else:
+            raise NotImplementedError(
+                f"Only isinstance for byte, int, Plutus Dataclass types are supported"
+            )
+
+
 class PythonBuiltIn(Enum):
     all = OLambda(
         ["xs"],
@@ -437,6 +511,7 @@ class PythonBuiltIn(Enum):
             OVar("xs"), plt.BuiltIn(uplc.BuiltInFun.AddInteger), plt.Integer(0)
         ),
     )
+    isinstance = "isinstance"
 
 
 PythonBuiltInTypes = {
@@ -510,4 +585,5 @@ class PythonBuiltIn(Enum):
             IntegerInstanceType,
         )
     ),
+    PythonBuiltIn.isinstance: InstanceType(PolymorphicFunctionType(IsinstanceImpl())),
 }