Add option to assemble constants

CHANGELOG.md

@@ -1,3 +1,9 @@
+# Future
+
+## Added
+* Added the `assembleConstants` option to `compileTeal`. When enabled, the compiler will assemble
+int and byte constants in the most efficient way to reduce program size ([#57](https://github.com/algorand/pyteal/pull/57)).
+
 # 0.7.0
 
 ## Added

mypy.ini

@@ -2,3 +2,6 @@
 
 [mypy-pytest.*]
 ignore_missing_imports = True
+
+[mypy-algosdk.*]
+ignore_missing_imports = True

pyteal/__init__.py

@@ -2,12 +2,15 @@
 from .ast import __all__ as ast_all
 from .ir import *
 from .ir import __all__ as ir_all
-from .compiler import CompileOptions, compileTeal
+from .compiler import MAX_TEAL_VERSION, MIN_TEAL_VERSION, DEFAULT_TEAL_VERSION, CompileOptions, compileTeal
 from .types import TealType
 from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
 from .config import MAX_GROUP_SIZE
 
 __all__ = ast_all + ir_all + [
+    "MAX_TEAL_VERSION",
+    "MIN_TEAL_VERSION",
+    "DEFAULT_TEAL_VERSION",
     "CompileOptions",
     "compileTeal",
     "TealType",

pyteal/compiler/__init__.py

@@ -0,0 +1,9 @@
+from .compiler import MAX_TEAL_VERSION, MIN_TEAL_VERSION, DEFAULT_TEAL_VERSION, CompileOptions, compileTeal
+
+__all__ = [
+    "MAX_TEAL_VERSION",
+    "MIN_TEAL_VERSION",
+    "DEFAULT_TEAL_VERSION",
+    "CompileOptions",
+    "compileTeal",
+]

pyteal/compiler.py → pyteal/compiler/compiler.py

@@ -1,10 +1,13 @@
-from typing import List, DefaultDict, cast
-from collections import defaultdict
+from typing import List
 
-from .ast import Expr
-from .ir import Op, Mode, TealComponent, TealOp, TealLabel, TealBlock, TealSimpleBlock, TealConditionalBlock
-from .errors import TealInputError, TealInternalError
-from .config import NUM_SLOTS
+from ..ast import Expr
+from ..ir import Mode, TealComponent, TealOp, TealBlock
+from ..errors import TealInputError, TealInternalError
+from ..config import NUM_SLOTS
+
+from .sort import sortBlocks
+from .flatten import flattenBlocks
+from .constants import createConstantBlocks
 
 MAX_TEAL_VERSION = 3
 MIN_TEAL_VERSION = 2
@@ -16,95 +19,6 @@ def __init__(self, *, mode: Mode = Mode.Signature, version: int = DEFAULT_TEAL_V
         self.mode = mode
         self.version = version
 
-def sortBlocks(start: TealBlock) -> List[TealBlock]:
-    """Topologically sort the graph which starts with the input TealBlock.
-
-    Args:
-        start: The starting point of the graph to sort.
-
-    Returns:
-        An ordered list of TealBlocks that is sorted such that every block is guaranteed to appear
-        in the list before all of its outgoing blocks.
-    """
-    # based on Kahn's algorithm from https://en.wikipedia.org/wiki/Topological_sorting
-    S = [start]
-    order = []
-
-    while len(S) != 0:
-        n = S.pop(0)
-        order.append(n)
-        for i, m in enumerate(n.getOutgoing()):
-            for i, block in enumerate(m.incoming):
-                if n is block:
-                    m.incoming.pop(i)
-                    break
-            if len(m.incoming) == 0:
-                if i == 0:
-                    S.insert(0, m)
-                else:
-                    S.append(m)
-
-    return order
-
-def flattenBlocks(blocks: List[TealBlock]) -> List[TealComponent]:
-    """Lowers a list of TealBlocks into a list of TealComponents.
-
-    Args:
-        blocks: The blocks to lower.
-    """
-    codeblocks = []
-    references: DefaultDict[int, int] = defaultdict(int)
-
-    indexToLabel = lambda index: "l{}".format(index)
-
-    for i, block in enumerate(blocks):
-        code = list(block.ops)
-        codeblocks.append(code)
-        if block.isTerminal():
-            continue
-
-        if type(block) is TealSimpleBlock:
-            simpleBlock = cast(TealSimpleBlock, block)
-            assert simpleBlock.nextBlock is not None
-
-            nextIndex = blocks.index(simpleBlock.nextBlock, i+1)
-            if nextIndex != i + 1:
-                references[nextIndex] += 1
-                code.append(TealOp(None, Op.b, indexToLabel(nextIndex)))
-        elif type(block) is TealConditionalBlock:
-            conditionalBlock = cast(TealConditionalBlock, block)
-            assert conditionalBlock.trueBlock is not None
-            assert conditionalBlock.falseBlock is not None
-
-            trueIndex = blocks.index(conditionalBlock.trueBlock, i+1)
-            falseIndex = blocks.index(conditionalBlock.falseBlock, i+1)
-
-            if falseIndex == i + 1:
-                references[trueIndex] += 1
-                code.append(TealOp(None, Op.bnz, indexToLabel(trueIndex)))
-                continue
-
-            if trueIndex == i + 1:
-                references[falseIndex] += 1
-                code.append(TealOp(None, Op.bz, indexToLabel(falseIndex)))
-                continue
-
-            references[trueIndex] += 1
-            code.append(TealOp(None, Op.bnz, indexToLabel(trueIndex)))
-
-            references[falseIndex] += 1
-            code.append(TealOp(None, Op.b, indexToLabel(falseIndex)))
-        else:
-            raise TealInternalError("Unrecognized block type: {}".format(type(block)))
-
-    teal: List[TealComponent] = []
-    for i, code in enumerate(codeblocks):
-        if references[i] != 0:
-            teal.append(TealLabel(None, indexToLabel(i)))
-        teal += code
-
-    return teal
-
 def verifyOpsForVersion(teal: List[TealComponent], version: int):
     """Verify that all TEAL operations are allowed in the specified version.
 
@@ -137,7 +51,7 @@ def verifyOpsForMode(teal: List[TealComponent], mode: Mode):
             if not op.mode & mode:
                 raise TealInputError("Op not supported in {} mode: {}".format(mode.name, op))
 
-def compileTeal(ast: Expr, mode: Mode, *, version: int = DEFAULT_TEAL_VERSION) -> str:
+def compileTeal(ast: Expr, mode: Mode, *, version: int = DEFAULT_TEAL_VERSION, assembleConstants: bool = False) -> str:
     """Compile a PyTeal expression into TEAL assembly.
 
     Args:
@@ -146,12 +60,18 @@ def compileTeal(ast: Expr, mode: Mode, *, version: int = DEFAULT_TEAL_VERSION) -
         version (optional): The TEAL version used to assemble the program. This will determine which
             expressions and fields are able to be used in the program and how expressions compile to
             TEAL opcodes. Defaults to 2 if not included.
+        assembleConstants (optional): When true, the compiler will produce a program with fully
+            assembled constants, rather than using the pseudo-ops `int`, `byte`, and `addr`. These
+            constants will be assembled in the most space-efficient way, so enabling this may reduce
+            the compiled program's size. Enabling this option requires a minimum TEAL version of 3.
+            Defaults to false.
 
     Returns:
         A TEAL assembly program compiled from the input expression.
 
     Raises:
         TealInputError: if an operation in ast is not supported by the supplied mode and version.
+        TealInternalError: if an internal error is encounter during compilation.
     """
     if not (MIN_TEAL_VERSION <= version <= MAX_TEAL_VERSION) or type(version) != int:
         raise TealInputError("Unsupported TEAL version: {}. Excepted an integer in the range [{}, {}]".format(version, MIN_TEAL_VERSION, MAX_TEAL_VERSION))
@@ -188,6 +108,11 @@ def compileTeal(ast: Expr, mode: Mode, *, version: int = DEFAULT_TEAL_VERSION) -
     for index, slot in enumerate(sorted(slots, key=lambda slot: slot.id)):
         for stmt in teal:
             stmt.assignSlot(slot, index)
+
+    if assembleConstants:
+        if version < 3:
+            raise TealInternalError("The minimum TEAL version required to enable assembleConstants is 3. The current version is {}".format(version))
+        teal = createConstantBlocks(teal)
 
     lines = ["#pragma version {}".format(version)]
     lines += [i.assemble() for i in teal]

pyteal/compiler/compiler_test.py

@@ -0,0 +1,174 @@
+import pytest
+
+from .. import *
+
+def test_compile_single():
+    expr = Int(1)
+
+    expected = """
+#pragma version 2
+int 1
+""".strip()
+    actual_application = compileTeal(expr, Mode.Application)
+    actual_signature = compileTeal(expr, Mode.Signature)
+
+    assert actual_application == actual_signature
+    assert actual_application == expected
+
+def test_compile_sequence():
+    expr = Seq([Pop(Int(1)), Pop(Int(2)), Int(3) + Int(4)])
+
+    expected = """
+#pragma version 2
+int 1
+pop
+int 2
+pop
+int 3
+int 4
++
+""".strip()
+    actual_application = compileTeal(expr, Mode.Application)
+    actual_signature = compileTeal(expr, Mode.Signature)
+
+    assert actual_application == actual_signature
+    assert actual_application == expected
+
+def test_compile_branch():
+    expr = If(Int(1), Bytes("true"), Bytes("false"))
+
+    expected = """
+#pragma version 2
+int 1
+bnz l2
+byte "false"
+b l3
+l2:
+byte "true"
+l3:
+""".strip()
+    actual_application = compileTeal(expr, Mode.Application)
+    actual_signature = compileTeal(expr, Mode.Signature)
+
+    assert actual_application == actual_signature
+    assert actual_application == expected
+
+def test_compile_mode():
+    expr = App.globalGet(Bytes("key"))
+
+    expected = """
+#pragma version 2
+byte "key"
+app_global_get
+""".strip()
+    actual_application = compileTeal(expr, Mode.Application)
+
+    assert actual_application == expected
+
+    with pytest.raises(TealInputError):
+        compileTeal(expr, Mode.Signature)
+
+def test_compile_version_invalid():
+    expr = Int(1)
+
+    with pytest.raises(TealInputError):
+        compileTeal(expr, Mode.Signature, version=1) # too small
+
+    with pytest.raises(TealInputError):
+        compileTeal(expr, Mode.Signature, version=4) # too large
+
+    with pytest.raises(TealInputError):
+        compileTeal(expr, Mode.Signature, version=2.0) # decimal
+
+def test_compile_version_2():
+    expr = Int(1)
+
+    expected = """
+#pragma version 2
+int 1
+""".strip()
+    actual = compileTeal(expr, Mode.Signature, version=2)
+    assert actual == expected
+
+def test_compile_version_default():
+    expr = Int(1)
+
+    actual_default = compileTeal(expr, Mode.Signature)
+    actual_version_2 = compileTeal(expr, Mode.Signature, version=2)
+    assert actual_default == actual_version_2
+
+def test_compile_version_3():
+    expr = Int(1)
+
+    expected = """
+#pragma version 3
+int 1
+""".strip()
+    actual = compileTeal(expr, Mode.Signature, version=3)
+    assert actual == expected
+
+def test_slot_load_before_store():
+
+    program = AssetHolding.balance(Int(0), Int(0)).value()
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2)
+
+    program = AssetHolding.balance(Int(0), Int(0)).hasValue()
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2)
+
+    program = App.globalGetEx(Int(0), Bytes("key")).value()
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2)
+
+    program = App.globalGetEx(Int(0), Bytes("key")).hasValue()
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2)
+
+    program = ScratchVar().load()
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2)
+
+def test_assembleConstants():
+    program = Itob(Int(1) + Int(1) + Int(2)) == Concat(Bytes("test"), Bytes("test"), Bytes("test2"))
+
+    expectedNoAssemble = """
+#pragma version 3
+int 1
+int 1
++
+int 2
++
+itob
+byte "test"
+byte "test"
+concat
+byte "test2"
+concat
+==
+""".strip()
+    actualNoAssemble = compileTeal(program, Mode.Application, version=3, assembleConstants=False)
+    assert expectedNoAssemble == actualNoAssemble
+
+    expectedAssemble = """
+#pragma version 3
+intcblock 1
+bytecblock 0x74657374
+intc_0 // 1
+intc_0 // 1
++
+pushint 2 // 2
++
+itob
+bytec_0 // "test"
+bytec_0 // "test"
+concat
+pushbytes 0x7465737432 // "test2"
+concat
+==
+""".strip()
+    actualAssemble = compileTeal(program, Mode.Application, version=3, assembleConstants=True)
+    assert expectedAssemble == actualAssemble
+
+    with pytest.raises(TealInternalError):
+        compileTeal(program, Mode.Application, version=2, assembleConstants=True)