phasm/compiler.py

@@ -8,6 +8,7 @@ from . import typing
 from . import wasm
 from . import wasmeasy
 
+from .stdlib import alloc as stdlib_alloc
 from .wasmeasy import i32, i64
 
 Statements = Generator[wasm.Statement, None, None]
@@ -362,10 +363,10 @@ def module(inp: ourlang.Module) -> wasm.Module:
     ]
 
     result.functions = [
-        _generate____new_reference___(inp),
+        _generate____new_reference___(inp), # Old allocator
-        _generate_stdlib_alloc___init__(inp),
+        stdlib_alloc.__init__,
-        _generate_stdlib_alloc___find_free_block__(inp),
+        stdlib_alloc.__find_free_block__,
-        _generate_stdlib_alloc___alloc__(inp),
+        stdlib_alloc.__alloc__,
         _generate____access_bytes_index___(inp),
     ] + [
         function(x)
@@ -398,127 +399,6 @@ def _generate____new_reference___(mod: ourlang.Module) -> wasm.Function:
         ],
     )
 
-STDLIB_ALLOC__IDENTIFIER = 0xA1C0
-STDLIB_ALLOC__RESERVED0 = 0x04
-STDLIB_ALLOC__FREE_BLOCK = 0x08
-STDLIB_ALLOC__UNALLOC_ADDR = 0x0C
-
-def _generate_stdlib_alloc___init__(mod: ourlang.Module) -> wasm.Function:
-    return wasm.Function(
-        'stdlib.alloc.__init__',
-        'stdlib.alloc.__init__',
-        [],
-        [],
-        wasm.WasmTypeNone(),
-        [
-            i32.const(0),
-            i32.load(),
-            i32.const(STDLIB_ALLOC__IDENTIFIER),
-            i32.eq(),
-            *wasmeasy.if_(
-                wasm.Statement('return', comment='Already set up'),
-            ),
-
-            i32.const(STDLIB_ALLOC__RESERVED0, comment='Reserved'),
-            i32.const(0),
-            i32.store(),
-
-            i32.const(STDLIB_ALLOC__FREE_BLOCK,
-                comment='Address of next free block'),
-            i32.const(0, comment='None to start with'),
-            i32.store(),
-
-            i32.const(STDLIB_ALLOC__UNALLOC_ADDR,
-                comment='Address of first unallocated byte'),
-            i32.const(0x10),
-            i32.store(),
-
-            i32.const(0, comment='Done setting up'),
-            i32.const(STDLIB_ALLOC__IDENTIFIER),
-            i32.store(),
-        ],
-    )
-
-def _generate_stdlib_alloc___find_free_block__(mod: ourlang.Module) -> wasm.Function:
-    return wasm.Function(
-        'stdlib.alloc.__find_free_block__',
-        'stdlib.alloc.__find_free_block__',
-        [
-            ('alloc_size', type_(mod.types['i32']), ),
-        ],
-        [
-            ('result', type_(mod.types['i32']), ),
-        ],
-        type_(mod.types['i32']),
-        [
-            i32.const(STDLIB_ALLOC__FREE_BLOCK),
-            i32.load(),
-            i32.const(0),
-            i32.eq(),
-            *wasmeasy.if_(
-                i32.const(0),
-                wasm.Statement('return'),
-            ),
-            wasm.Statement('unreachable'),
-        ],
-    )
-
-def _generate_stdlib_alloc___alloc__(mod: ourlang.Module) -> wasm.Function:
-    return wasm.Function(
-        'stdlib.alloc.__alloc__',
-        'stdlib.alloc.__alloc__',
-        [
-            ('alloc_size', type_(mod.types['i32']), ),
-        ],
-        [
-            ('result', type_(mod.types['i32']), ),
-        ],
-        type_(mod.types['i32']),
-        [
-            i32.const(0),
-            i32.load(),
-            i32.const(STDLIB_ALLOC__IDENTIFIER),
-            i32.ne(),
-            *wasmeasy.if_(
-                wasm.Statement('unreachable'),
-            ),
-
-            wasm.Statement('local.get', '$alloc_size'),
-            wasm.Statement('call', '$stdlib.alloc.__find_free_block__'),
-            wasm.Statement('local.set', '$result'),
-
-            # Check if there was a free block
-            wasm.Statement('local.get', '$result'),
-            i32.const(0),
-            i32.eq(),
-            *wasmeasy.if_(
-                # Use unallocated space
-                i32.const(STDLIB_ALLOC__UNALLOC_ADDR),
-
-                i32.const(STDLIB_ALLOC__UNALLOC_ADDR),
-                i32.load(),
-                wasm.Statement('local.tee', '$result'),
-
-                # Updated unalloc pointer (address already set on stack)
-                i32.const(4), # Header size
-                i32.add(),
-                wasm.Statement('local.get', '$alloc_size'),
-                i32.add(),
-                i32.store('offset=0'),
-            ),
-
-            # Store block size
-            wasm.Statement('local.get', '$result'),
-            wasm.Statement('local.get', '$alloc_size'),
-            i32.store('offset=0'),
-
-            # Return address of the allocated bytes
-            wasm.Statement('local.get', '$result'),
-            i32.const(4), # Header size
-            i32.add(),
-        ],
-    )
-
 def _generate____access_bytes_index___(mod: ourlang.Module) -> wasm.Function:
     return wasm.Function(
         '___access_bytes_index___',

phasm/stdlib/alloc.py

@@ -1,7 +1,7 @@
 """
 stdlib: Memory allocation
 """
-from phasm.wasmeasy import Generator, func_wrapper
+from phasm.wasmgenerator import Generator, VarType_i32 as i32, func_wrapper
 
 IDENTIFIER = 0xA1C0
 
@@ -10,18 +10,19 @@ ADR_RESERVED0 = ADR_IDENTIFIER + 4
 ADR_FREE_BLOCK_PTR = ADR_RESERVED0 + 4
 ADR_UNALLOC_PTR = ADR_FREE_BLOCK_PTR + 4
 
-@func_wrapper
+UNALLOC_PTR = ADR_UNALLOC_PTR + 4
+
+@func_wrapper()
 def __init__(g: Generator) -> None:
     """
     Initializes the memory so we can allocate it
     """
 
     # Check if the memory is already initialized
-    g.i32.const(0)
+    g.i32.const(ADR_IDENTIFIER)
     g.i32.load()
     g.i32.const(IDENTIFIER)
     g.i32.eq()
-
     with g.if_():
         # Already initialized, return without any changes
         g.return_()
@@ -40,7 +41,7 @@ def __init__(g: Generator) -> None:
     # Store the pointer towards the first unallocated block
     # In this case the end of the stdlib.alloc header at the start
     g.i32.const(ADR_UNALLOC_PTR)
-    g.i32.const(0x10)
+    g.i32.const(UNALLOC_PTR)
     g.i32.store()
 
     # Store that we've initialized the memory
@@ -48,9 +49,8 @@ def __init__(g: Generator) -> None:
     g.i32.const(IDENTIFIER)
     g.i32.store()
 
-
+@func_wrapper(exported=False)
-@func_wrapper
+def __find_free_block__(g: Generator, alloc_size: i32) -> i32:
-def __find_free_block__(g: Generator) -> None:
     # Find out if we've freed any blocks at all so far
     g.i32.const(ADR_FREE_BLOCK_PTR)
     g.i32.load()
@@ -61,69 +61,61 @@ def __find_free_block__(g: Generator) -> None:
         g.i32.const(0)
         g.return_()
 
+    del alloc_size # TODO
     g.unreachable()
 
-@func_wrapper
+    return i32('return') # To satisfy mypy
-def __alloc__(g: Generator) -> None:
+
+@func_wrapper()
+def __alloc__(g: Generator, alloc_size: i32) -> i32:
+    result = i32('result')
+
+    # Check if the memory is already initialized
     g.i32.const(ADR_IDENTIFIER)
     g.i32.load()
     g.i32.const(IDENTIFIER)
     g.i32.ne()
     with g.if_():
+        # Not yet initialized, or memory corruption
         g.unreachable()
 
-# def _generate_stdlib_alloc___alloc__(mod: ourlang.Module) -> wasm.Function:
+    # Try to claim a free block
-#     return wasm.Function(
+    g.local.get(alloc_size)
-#         'stdlib.alloc.__alloc__',
+    g.call(__find_free_block__)
-#         'stdlib.alloc.__alloc__',
+    g.local.set(result)
-#         [
+
-#             ('alloc_size', type_(mod.types['i32']), ),
+    # Check if there was a free block
-#         ],
+    g.local.get(result)
-#         [
+    g.i32.const(0)
-#             ('result', type_(mod.types['i32']), ),
+    g.i32.eq()
-#         ],
+    with g.if_():
-#         type_(mod.types['i32']),
+        # No free blocks, increase allocated memory usage
-#         [
+
-#             i32.const(0),
+        # Put the address on the stack in advance so we can store to it later
-#             i32.load(),
+        g.i32.const(ADR_UNALLOC_PTR)
-#             i32.const(STDLIB_ALLOC__IDENTIFIER),
+
-#             i32.ne(),
+        # Get the current unalloc pointer value
-#             *wasmeasy.if_(
+        g.i32.const(ADR_UNALLOC_PTR)
-#                 wasm.Statement('unreachable'),
+        g.i32.load()
-#             ),
+        g.local.tee(result)
-#
+
-#             wasm.Statement('local.get', '$alloc_size'),
+        # Calculate new unalloc pointer value
-#             wasm.Statement('call', '$stdlib.alloc.__find_free_block__'),
+        g.i32.const(4) # Header size
-#             wasm.Statement('local.set', '$result'),
+        g.i32.add()
-#
+        g.local.get(alloc_size)
-#             # Check if there was a free block
+        g.i32.add()
-#             wasm.Statement('local.get', '$result'),
+
-#             i32.const(0),
+        # Store new unalloc pointer value (address was set on stack in advance)
-#             i32.eq(),
+        g.i32.store()
-#             *wasmeasy.if_(
+
-#                 # Use unallocated space
+    # Store block size in the header
-#                 i32.const(STDLIB_ALLOC__ADR_UNALLOC_PTR),
+    g.local.get(result)
-#
+    g.local.get(alloc_size)
-#                 i32.const(STDLIB_ALLOC__ADR_UNALLOC_PTR),
+    g.i32.store()
-#                 i32.load(),
+
-#                 wasm.Statement('local.tee', '$result'),
+    # Return address of the allocated bytes
-#
+    g.local.get(result)
-#                 # Updated unalloc pointer (address already set on stack)
+    g.i32.const(4) # Header size
-#                 i32.const(4), # Header size
+    g.i32.add()
-#                 i32.add(),
+
-#                 wasm.Statement('local.get', '$alloc_size'),
+    return i32('return') # To satisfy mypy
-#                 i32.add(),
-#                 i32.store('offset=0'),
-#             ),
-#
-#             # Store block size
-#             wasm.Statement('local.get', '$result'),
-#             wasm.Statement('local.get', '$alloc_size'),
-#             i32.store('offset=0'),
-#
-#             # Return address of the allocated bytes
-#             wasm.Statement('local.get', '$result'),
-#             i32.const(4), # Header size
-#             i32.add(),
-#         ],
-#     )

phasm/wasmeasy.py

@@ -1,7 +1,7 @@
 """
 Helper functions to quickly generate WASM code
 """
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional, Type
 
 import functools
 
@@ -67,76 +67,3 @@ class Block:
         ]
 
 if_ = Block('if')
-
-class Generator_i32:
-    def __init__(self, generator: 'Generator') -> None:
-        self.generator = generator
-
-        # 2.4.1. Numeric Instructions
-        self.eq = functools.partial(self.generator.add, 'i32.eq')
-        self.ne = functools.partial(self.generator.add, 'i32.ne')
-
-        # 2.4.4. Memory Instructions
-        self.load = functools.partial(self.generator.add, 'i32.load')
-        self.store = functools.partial(self.generator.add, 'i32.store')
-
-    def const(self, value: int, comment: Optional[str] = None) -> None:
-        self.generator.add('i32.const', f'0x{value:08x}', comment=comment)
-
-class GeneratorBlock:
-    def __init__(self, generator: 'Generator', name: str) -> None:
-        self.generator = generator
-        self.name = name
-
-    def __enter__(self) -> None:
-        self.generator.add(self.name)
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        if not exc_type:
-            self.generator.add('end')
-
-class Generator:
-    def __init__(self) -> None:
-        self.statements: List[wasm.Statement] = []
-
-        self.i32 = Generator_i32(self)
-
-        # 2.4.5 Control Instructions
-        self.nop = functools.partial(self.add, 'nop')
-        self.unreachable = functools.partial(self.add, 'unreachable')
-        # block
-        # loop
-        self.if_ = functools.partial(GeneratorBlock, self, 'if')
-        # br
-        # br_if
-        # br_table
-        self.return_ = functools.partial(self.add, 'return')
-        # call
-        # call_indirect
-
-    def add(self, name: str, *args: str, comment: Optional[str] = None) -> None:
-        self.statements.append(wasm.Statement(name, *args, comment=comment))
-
-def func_wrapper(func: Any) -> wasm.Function:
-    assert func.__module__.startswith('phasm.')
-
-    assert Generator is func.__annotations__['g']
-
-    params: Any = []
-
-    locals_: Any = []
-
-    assert None is func.__annotations__['return']
-    return_type = wasm.WasmTypeNone()
-
-    generator = Generator()
-    func(g=generator)
-
-    return wasm.Function(
-        func.__module__[6:] + '.' + func.__name__,
-        func.__module__[6:] + '.' + func.__name__,
-        params,
-        locals_,
-        return_type,
-        generator.statements,
-    )

phasm/wasmgenerator.py

@@ -0,0 +1,154 @@
+"""
+Helper functions to generate WASM code by writing Python functions
+"""
+from typing import Any, Callable, Dict, List, Optional, Type
+
+import functools
+
+from . import wasm
+
+# pylint: disable=C0103,C0115,C0116,R0902
+
+class VarType_Base:
+    wasm_type: Type[wasm.WasmType]
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.name_ref = f'${name}'
+
+class VarType_i32(VarType_Base):
+    wasm_type = wasm.WasmTypeInt32
+
+class Generator_i32:
+    def __init__(self, generator: 'Generator') -> None:
+        self.generator = generator
+
+        # 2.4.1. Numeric Instructions
+        # ibinop
+        self.add = functools.partial(self.generator.add, 'i32.add')
+
+        # irelop
+        self.eq = functools.partial(self.generator.add, 'i32.eq')
+        self.ne = functools.partial(self.generator.add, 'i32.ne')
+
+        # 2.4.4. Memory Instructions
+        self.load = functools.partial(self.generator.add, 'i32.load')
+        self.store = functools.partial(self.generator.add, 'i32.store')
+
+    def const(self, value: int, comment: Optional[str] = None) -> None:
+        self.generator.add('i32.const', f'0x{value:08x}', comment=comment)
+
+class Generator_Local:
+    def __init__(self, generator: 'Generator') -> None:
+        self.generator = generator
+
+    # 2.4.3. Variable Instructions
+    def get(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
+        self.generator.add('local.get', variable.name_ref, comment=comment)
+
+    def set(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
+        self.generator.locals.setdefault(variable.name, variable)
+
+        self.generator.add('local.set', variable.name_ref, comment=comment)
+
+    def tee(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
+        self.generator.locals.setdefault(variable.name, variable)
+
+        self.generator.add('local.tee', variable.name_ref, comment=comment)
+
+class GeneratorBlock:
+    def __init__(self, generator: 'Generator', name: str) -> None:
+        self.generator = generator
+        self.name = name
+
+    def __enter__(self) -> None:
+        self.generator.add(self.name)
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        if not exc_type:
+            self.generator.add('end')
+
+class Generator:
+    def __init__(self) -> None:
+        self.statements: List[wasm.Statement] = []
+        self.locals: Dict[str, VarType_Base] = {}
+
+        self.i32 = Generator_i32(self)
+
+        # 2.4.3 Variable Instructions
+        self.local = Generator_Local(self)
+
+        # 2.4.5 Control Instructions
+        self.nop = functools.partial(self.add, 'nop')
+        self.unreachable = functools.partial(self.add, 'unreachable')
+        # block
+        # loop
+        self.if_ = functools.partial(GeneratorBlock, self, 'if')
+        # br
+        # br_if
+        # br_table
+        self.return_ = functools.partial(self.add, 'return')
+        # call
+        # call_indirect
+
+    def call(self, function: wasm.Function) -> None:
+        self.add('call', f'${function.name}')
+
+    def add(self, name: str, *args: str, comment: Optional[str] = None) -> None:
+        self.statements.append(wasm.Statement(name, *args, comment=comment))
+
+def func_wrapper(exported: bool = True) -> Callable[[Any], wasm.Function]:
+    """
+    This wrapper will execute the function and return
+    a wasm Function with the generated Statements
+    """
+    def inner(func: Any) -> wasm.Function:
+        func_name_parts = func.__module__.split('.') + [func.__name__]
+        if 'phasm' == func_name_parts[0]:
+            func_name_parts.pop(0)
+        func_name = '.'.join(func_name_parts)
+
+        annot = dict(func.__annotations__)
+
+        # Check if we can pass the generator
+        assert Generator is annot.pop('g')
+
+        # Convert return type to WasmType
+        annot_return = annot.pop('return')
+        if annot_return is None:
+            return_type = wasm.WasmTypeNone()
+        else:
+            assert issubclass(annot_return, VarType_Base)
+            return_type = annot_return.wasm_type()
+
+        # Load the argument types, and generate instances
+        args: Dict[str, VarType_Base] = {}
+        params: List[wasm.Param] = []
+        for param_name, param_type in annot.items():
+            assert issubclass(param_type, VarType_Base)
+
+            params.append((param_name, param_type.wasm_type(), ))
+
+            args[param_name] = VarType_Base(param_name)
+
+        # Make a generator, and run the function on that generator,
+        # so the Statements get added
+        generator = Generator()
+        func(g=generator, **args)
+
+        # Check what locals were used, and define them
+        locals_: List[wasm.Param] = []
+        for local_name, local_type in generator.locals.items():
+            locals_.append((local_name, local_type.wasm_type(), ))
+
+        # Complete function definition
+        return wasm.Function(
+            func_name,
+            func_name if exported else None,
+            params,
+            locals_,
+            return_type,
+            generator.statements,
+        )
+
+    return inner