From 9616d2046073c62692cc9ba389bccb34786c8f7c Mon Sep 17 00:00:00 2001
From: "Johan B.W. de Vries" <info@jbwdevries.nl>
Date: Sat, 7 Aug 2021 15:02:20 +0200
Subject: [PATCH] fib works \o/

---
 py2wasm/python.py             | 44 +++++++++++++++++++++++++++++++----
 tests/integration/test_fib.py |  2 +-
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/py2wasm/python.py b/py2wasm/python.py
index 9c6e989..57d004e 100644
--- a/py2wasm/python.py
+++ b/py2wasm/python.py
@@ -29,21 +29,36 @@ class Visitor:
 
         module = wasm.Module()
 
+        # Do a check first for all function definitions
+        # to get their types. Otherwise you cannot call
+        # a method that you haven't defined just yet,
+        # even if it is in the same file
+        function_body_map: Dict[ast.FunctionDef, wasm.Function] = {}
         for stmt in node.body:
             if isinstance(stmt, ast.FunctionDef):
-                wnode = self.visit_FunctionDef(module, stmt)
+                wnode = self.pre_visit_FunctionDef(module, stmt)
 
                 if isinstance(wnode, wasm.Import):
                     module.imports.append(wnode)
                 else:
                     module.functions.append(wnode)
+                    function_body_map[stmt] = wnode
+                continue
+
+            # No other pre visits to do
+
+        for stmt in node.body:
+            if isinstance(stmt, ast.FunctionDef):
+                if stmt in function_body_map:
+                    self.parse_FunctionDef_body(module, function_body_map[stmt], stmt)
+                # else: It's an import, no actual body to parse
                 continue
 
             raise NotImplementedError(stmt)
 
         return module
 
-    def visit_FunctionDef(
+    def pre_visit_FunctionDef(
             self,
             module: wasm.Module,
             node: ast.FunctionDef,
@@ -55,6 +70,8 @@ class Visitor:
 
         Nested / dynamicly created functions are not yet supported
         """
+        del module
+
         exported = False
 
         if node.decorator_list:
@@ -99,7 +116,18 @@ class Visitor:
             ]
         ]
 
-        wlocals: WLocals = dict(params)
+        return wasm.Function(node.name, exported, params, result, [])
+
+    def parse_FunctionDef_body(
+            self,
+            module: wasm.Module,
+            func: wasm.Function,
+            node: ast.FunctionDef,
+        ) -> None:
+        """
+        Parses the function body
+        """
+        wlocals: WLocals = dict(func.params)
 
         statements: List[wasm.Statement] = []
         for py_stmt in node.body:
@@ -107,7 +135,7 @@ class Visitor:
                 self.visit_stmt(module, node, wlocals, py_stmt)
             )
 
-        return wasm.Function(node.name, exported, params, result, statements)
+        func.statements = statements
 
     def visit_stmt(
             self,
@@ -250,6 +278,10 @@ class Visitor:
             yield wasm.Statement('{}.add'.format(exp_type))
             return
 
+        if isinstance(node.op, ast.Sub):
+            yield wasm.Statement('{}.sub'.format(exp_type))
+            return
+
         raise NotImplementedError(node.op)
 
     def visit_Compare(
@@ -274,6 +306,10 @@ class Visitor:
             yield wasm.Statement('i32.lt_s')
             return
 
+        if isinstance(node.ops[0], ast.Eq):
+            yield wasm.Statement('i32.eq')
+            return
+
         if isinstance(node.ops[0], ast.Gt):
             yield wasm.Statement('i32.gt_s')
             return
diff --git a/tests/integration/test_fib.py b/tests/integration/test_fib.py
index a68cb0e..b27297d 100644
--- a/tests/integration/test_fib.py
+++ b/tests/integration/test_fib.py
@@ -21,7 +21,7 @@ def fib(n: i32) -> i32:
     return helper(n - 1, 0, 1)
 
 @exported
-def testEntry():
+def testEntry() -> i32:
     return fib(40)
 """