Richer types

examples/rich_types.tl

@@ -0,0 +1,9 @@
+#pragma version 8
+
+const int I = 123
+const bytes B = "abc"
+const bigint Bi = 1231231231231231231231231231231231 
+const bigint One = 1
+assert((Bi b/ Bi) b== One)
+
+exit(1)

tealish/base.py

@@ -1,6 +1,13 @@
 from typing import cast, Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 from tealish.errors import CompileError
-from .tealish_builtins import AVMType, ConstValue, ScratchRecord, VarType
+from .tealish_builtins import (
+    AVMType,
+    ConstValue,
+    ScratchRecord,
+    VarType,
+    TealishType,
+    stack_type,
+)
 from .langspec import get_active_langspec, Op
 from .scope import Scope
 
@@ -15,17 +22,19 @@
 def check_arg_types(name: str, incoming_args: List["Node"]) -> None:
     op = lang_spec.lookup_op(name)
     expected_args = op.arg_types
-    # TODO:
     for i, incoming_arg in enumerate(incoming_args):
-        if incoming_arg.type == AVMType.any:  # type: ignore
+        tealish_type = incoming_arg.tealish_type()
+        avm_type = stack_type(tealish_type)
+
+        if avm_type == AVMType.any:
             continue
         if expected_args[i] == AVMType.any:
             continue
-        if incoming_arg.type == expected_args[i]:  # type: ignore
+        if tealish_type == expected_args[i]:
             continue
 
         raise Exception(
-            f"Incorrect type {incoming_arg.type} "  # type: ignore
+            f"Incorrect type {tealish_type} "  # type: ignore
             + f"for arg {i} of {name}. Expected {expected_args[i]}"
         )
 
@@ -142,7 +151,7 @@ def check_arg_types(self, name: str, args: List["Node"]) -> None:
         except Exception as e:
             raise CompileError(str(e), node=self)  # type: ignore
 
-    def get_field_type(self, namespace: str, name: str) -> AVMType:
+    def get_field_type(self, namespace: str, name: str) -> TealishType:
         return lang_spec.get_field_type(namespace, name)
 
     def lookup_op(self, name: str) -> Op:
@@ -154,12 +163,20 @@ def lookup_func(self, name: str) -> "Func":
     def lookup_var(self, name: str) -> Any:
         return self.get_scope().lookup_var(name)
 
-    def lookup_const(self, name: str) -> Tuple["AVMType", ConstValue]:
+    def lookup_const(self, name: str) -> Tuple["TealishType", ConstValue]:
         return self.get_scope().lookup_const(name)
 
-    def lookup_avm_constant(self, name: str) -> Tuple["AVMType", Any]:
+    def lookup_avm_constant(self, name: str) -> Tuple["TealishType", Any]:
         return lang_spec.lookup_avm_constant(name)
 
+    def tealish_type(self) -> TealishType:
+        if hasattr(self, "type"):
+            return cast(TealishType, getattr(self, "type"))
+        return TealishType.none
+
+    def stack_type(self) -> AVMType:
+        return stack_type(self.tealish_type())
+
     # TODO: these attributes are only available on Node and other children types
     # we should either define them here or something else?
     @property

tealish/expression_nodes.py

@@ -2,7 +2,13 @@
 
 from .base import BaseNode
 from .errors import CompileError
-from .tealish_builtins import AVMType, get_struct, ObjectType
+from .tealish_builtins import (
+    AVMType,
+    get_struct,
+    ObjectType,
+    TealishType,
+    stack_type,
+)
 from .langspec import Op, type_lookup
 
 
@@ -14,7 +20,7 @@
 class Integer(BaseNode):
     def __init__(self, value: str, parent: Optional[BaseNode] = None) -> None:
         self.value = int(value)
-        self.type = AVMType.int
+        self.type = TealishType.int
         self.parent = parent
 
     def write_teal(self, writer: "TealWriter") -> None:
@@ -27,7 +33,7 @@ def _tealish(self) -> str:
 class Bytes(BaseNode):
     def __init__(self, value: str, parent: Optional[BaseNode] = None) -> None:
         self.value = value
-        self.type = AVMType.bytes
+        self.type = TealishType.bytes
         self.parent = parent
 
     def write_teal(self, writer: "TealWriter") -> None:
@@ -49,8 +55,8 @@ def process(self) -> None:
             raise CompileError(e.args[0], node=self)
         # is it a struct or box?
         if type(self.type) == tuple:
-            if self.type[0] == ObjectType.struct:
-                self.type = AVMType.bytes
+            if self.type[0] == ObjectType.scratch:
+                self.type = TealishType.bytes
             elif self.type[0] == ObjectType.box:
                 raise CompileError("Invalid use of a Box reference", node=self)
 
@@ -64,8 +70,8 @@ def _tealish(self) -> str:
 class Constant(BaseNode):
     def __init__(self, name: str, parent: Optional[BaseNode] = None) -> None:
         self.name = name
-        self.type: AVMType = AVMType.none
         self.parent = parent
+        self.type: TealishType
 
     def process(self) -> None:
         type, value = None, None
@@ -80,16 +86,16 @@ def process(self) -> None:
                 raise CompileError(
                     f'Constant "{self.name}" not declared in scope', node=self
                 )
-        if type not in (AVMType.int, AVMType.bytes):
+        if stack_type(type) not in (AVMType.int, AVMType.bytes):
             raise CompileError(f"Unexpected const type {type}", node=self)
 
         self.type = type
         self.value = value
 
     def write_teal(self, writer: "TealWriter") -> None:
-        if self.type == AVMType.int:
+        if stack_type(self.type) == AVMType.int:
             writer.write(self, f"pushint {self.value} // {self.name}")  # type: ignore
-        elif self.type == AVMType.bytes:
+        elif stack_type(self.type) == AVMType.bytes:
             writer.write(self, f"pushbytes {self.value} // {self.name}")  # type: ignore
 
     def _tealish(self) -> str:
@@ -153,7 +159,7 @@ def __init__(
 
     def process(self) -> None:
         self.expression.process()
-        self.type = self.expression.type
+        self.type = self.expression.tealish_type()
 
     def write_teal(self, writer: "TealWriter") -> None:
         writer.write(self, self.expression)
@@ -169,7 +175,7 @@ def __init__(
         self.name = name
         self.args = args
         self.parent = parent
-        self.type: Union[AVMType, List[AVMType]] = AVMType.none
+        self.type: Union[TealishType, List[TealishType]] = TealishType.none
         self.func_call_type: str = ""
         self.nodes = args
         self.immediate_args = ""
@@ -235,7 +241,7 @@ def process_op_call(self, op: Op) -> None:
     def process_special_call(self) -> None:
         self.func_call_type = "special"
         if self.name == "pop":
-            self.type = AVMType.any
+            self.type = TealishType.any
         for arg in self.args:
             arg.process()
 
@@ -275,7 +281,7 @@ def _tealish(self) -> str:
 class TxnField(BaseNode):
     def __init__(self, field: str, parent: Optional[BaseNode] = None) -> None:
         self.field = field
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -297,7 +303,7 @@ def __init__(
     ) -> None:
         self.field = field
         self.arrayIndex = arrayIndex
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -324,7 +330,7 @@ def __init__(
     ) -> None:
         self.field = field
         self.index = index
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -360,7 +366,7 @@ def __init__(
         self.field = field
         self.index = index
         self.arrayIndex = arrayIndex
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -403,7 +409,7 @@ def _tealish(self) -> str:
 class PositiveGroupIndex(BaseNode):
     def __init__(self, index: int, parent: Optional["BaseNode"] = None) -> None:
         self.index = index
-        self.type = AVMType.int
+        self.type = TealishType.int
         self.parent = parent
 
     def write_teal(self, writer: "TealWriter") -> None:
@@ -418,7 +424,7 @@ def _tealish(self) -> str:
 class NegativeGroupIndex(BaseNode):
     def __init__(self, index: int, parent: Optional[BaseNode] = None) -> None:
         self.index = index
-        self.type = AVMType.int
+        self.type = TealishType.int
         self.parent = parent
 
     def write_teal(self, writer: "TealWriter") -> None:
@@ -433,7 +439,7 @@ def _tealish(self) -> str:
 class GlobalField(BaseNode):
     def __init__(self, field: str, parent: Optional[BaseNode] = None) -> None:
         self.field = field
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -449,7 +455,7 @@ def _tealish(self) -> str:
 class InnerTxnField(BaseNode):
     def __init__(self, field: str, parent: Optional[BaseNode] = None) -> None:
         self.field = field
-        self.type = AVMType.any
+        self.type = TealishType.any
         self.parent = parent
 
     def process(self) -> None:
@@ -466,7 +472,7 @@ class StructOrBoxField(BaseNode):
     def __init__(self, name, field, parent=None) -> None:
         self.name = name
         self.field = field
-        self.type = AVMType.none
+        self.type = TealishType.none
         self.parent = parent
 
     def process(self) -> None:
@@ -479,9 +485,9 @@ def process(self) -> None:
         self.type = struct_field.data_type
 
     def write_teal(self, writer: "TealWriter") -> None:
-        if self.object_type == ObjectType.struct:
+        if self.object_type == ObjectType.scratch:
             writer.write(self, f"load {self.slot} // {self.name}")
-            if self.type == AVMType.int:
+            if self.type == TealishType.int:
                 writer.write(self, f"pushint {self.offset}")
                 writer.write(self, f"extract_uint64 // {self.field}")
             else:
@@ -491,7 +497,7 @@ def write_teal(self, writer: "TealWriter") -> None:
             writer.write(self, f"pushint {self.offset} // offset")
             writer.write(self, f"pushint {self.size} // size")
             writer.write(self, f"box_extract // {self.name}.{self.field}")
-            if self.type == AVMType.int:
+            if self.type == TealishType.int:
                 writer.write(self, "btoi")
         else:
             raise Exception()

tealish/langspec.py

@@ -3,25 +3,25 @@
 import requests
 import tealish
 import json
-from .tealish_builtins import constants, AVMType
+from .tealish_builtins import constants, TealishType
 from typing import List, Dict, Any, Tuple, Optional
 
 abc = "ABCDEFGHIJK"
 
 
 _opcode_type_map = {
-    ".": AVMType.any,
-    "B": AVMType.bytes,
-    "U": AVMType.int,
-    "": AVMType.none,
+    ".": TealishType.any,
+    "B": TealishType.bytes,
+    "U": TealishType.int,
+    "": TealishType.none,
 }
 
 
-def type_lookup(a: str) -> AVMType:
+def type_lookup(a: str) -> TealishType:
     return _opcode_type_map[a]
 
 
-def convert_args_to_types(args: str) -> List[AVMType]:
+def convert_args_to_types(args: str) -> List[TealishType]:
     return [type_lookup(args[idx]) for idx in range(len(args))]
 
 
@@ -35,21 +35,21 @@ class Op:
     #: list of arg types this op takes off the stack, encoded as a string
     args: str
     #: decoded list of incoming args
-    arg_types: List[AVMType]
+    arg_types: List[TealishType]
     #: list of arg types this op puts on the stack, encoded as a string
     returns: str
     #: decoded list of outgoing args
-    returns_types: List[AVMType]
+    returns_types: List[TealishType]
     #: how many bytes this opcode takes up when assembled
     size: int
     #: describes the args to be passed as immediate arguments to this op
     immediate_note: str
     #: describes the list of names that can be used as immediate arguments
     arg_enum: List[str]
     #: describes the types returned when each arg enum is used
-    arg_enum_types: List[AVMType]
+    arg_enum_types: List[TealishType]
     #: dictionary mapping the names in arg_enum to types in arg_enum_types
-    arg_enum_dict: Dict[str, AVMType]
+    arg_enum_dict: Dict[str, TealishType]
 
     #: informational string about the op
     doc: str
@@ -91,7 +91,7 @@ def __init__(self, op_def: Dict[str, Any]):
             if "ArgEnumTypes" in op_def:
                 self.arg_enum_types = convert_args_to_types(op_def["ArgEnumTypes"])
             else:
-                self.arg_enum_types = [AVMType.int] * len(self.arg_enum)
+                self.arg_enum_types = [TealishType.int] * len(self.arg_enum)
             self.arg_enum_dict = dict(zip(self.arg_enum, self.arg_enum_types))
         else:
             self.arg_enum = []
@@ -126,8 +126,8 @@ def __init__(self, spec: Dict[str, Any]) -> None:
             "Txn": self.ops["txn"].arg_enum_dict,
         }
 
-        self.global_fields: Dict[str, AVMType] = self.fields["Global"]
-        self.txn_fields: Dict[str, AVMType] = self.fields["Txn"]
+        self.global_fields: Dict[str, TealishType] = self.fields["Global"]
+        self.txn_fields: Dict[str, TealishType] = self.fields["Txn"]
 
     def as_dict(self) -> Dict[str, Any]:
         return self.spec
@@ -141,12 +141,12 @@ def lookup_op(self, name: str) -> Op:
             raise KeyError(f'Op "{name}" does not exist!')
         return self.ops[name]
 
-    def lookup_avm_constant(self, name: str) -> Tuple[AVMType, Any]:
+    def lookup_avm_constant(self, name: str) -> Tuple[TealishType, Any]:
         if name not in constants:
             raise KeyError(f'Constant "{name}" does not exist!')
         return constants[name]
 
-    def get_field_type(self, namespace: str, name: str) -> AVMType:
+    def get_field_type(self, namespace: str, name: str) -> TealishType:
         if "txn" in namespace:
             return self.txn_fields[name]
         elif namespace == "global":