feat(typecheck): Start typechecking (add statements this commit)

Author: MarcellPerger1

parser/astgen/astgen.py

@@ -112,6 +112,7 @@ def _walk_smt(self, smt: AnyNode) -> list[AstNode]:
         elif isinstance(smt, ConditionalBlock):
             return self._walk_conditional(smt)
         elif isinstance(smt, AssignNode):  # Simple assignment
+            # TODO: maybe separate SetAttr, SetItem, SetVar nodes?
             return [AstAssign(smt.region, self._walk_assign_left(smt.target),
                               self._walk_expr(smt.source))]
         elif isinstance(smt, AssignOpNode):  # Other (aug.) assignment

parser/typecheck/typecheck.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
+from typing import TypeAlias
 
 from util.recursive_eq import recursive_eq
-from ..astgen.ast_nodes import (
-    AstNode, AstIdent, AstDeclNode, AstDefine, VarDeclType, VarDeclScope)
+from ..astgen.ast_nodes import *
 from ..astgen.astgen import AstGen
 from ..astgen.filtered_walker import FilteredWalker
-from ..common import BaseLocatedError, StrRegion
+from ..common import BaseLocatedError, StrRegion, region_union, RegionUnionArgT
 
 
 @dataclass
@@ -193,7 +194,18 @@ def err(self, msg: str, region: StrRegion):
         return NameResolutionError(msg, region, self.src)
 
 
+class TypecheckError(BaseLocatedError):
+    """Errors raised by the typechecker"""
+
+
+NodeTypecheckFn: TypeAlias = 'Callable[[Typechecker, AstNode], TypeInfo | None]'
+
+_typecheck_dispatch: dict[type[AstNode], NodeTypecheckFn] = {}
+
+
 class Typechecker:
+    _curr_scope: Scope
+
     def __init__(self, name_resolver: NameResolver):
         self.resolver = name_resolver
         self.src = self.resolver.src
@@ -203,16 +215,118 @@ def _init(self):
         self.resolver.run()
         self.ast = self.resolver.ast
         self.top_scope = self.resolver.top_scope
+        self._curr_scope = self.top_scope
 
     def run(self):
         if self.is_ok is None:
             return self.is_ok
-        self._typecheck()
-        self.is_ok = True
+        self._init()
+        self._typecheck(self.ast)
+        self.is_ok = True  # didn't raise any errors
         return self.is_ok
 
-    def _typecheck(self):
-        walker = FilteredWalker()
-
-        self.ast.walk(walker)
-        ...
+    def _node_typechecker(self, tp=None):
+        if tp is None:
+            assert callable(self)
+            tp = self  # Called as decor in this class
+
+        def decor(fn: NodeTypecheckFn):
+            _typecheck_dispatch[tp] = fn
+            return fn
+        return decor
+
+    def _typecheck(self, n: AstNode):
+        try:
+            fn = _typecheck_dispatch[type(n)]
+        except KeyError:
+            fn = type(self)._typecheck_node_fallback
+        return fn(self, n)
+
+    def _typecheck_node_fallback(self, n: AstNode):
+        raise NotImplementedError(f"No typechecker function for node "
+                                  f"type {type(n).__name__}")
+
+    @_node_typechecker(AstProgramNode)
+    def _typecheck_program(self, n: AstProgramNode):
+        self._typecheck_block(n.statements)
+
+    def _typecheck_block(self, block: list[AstNode]):
+        for smt in block:
+            if (tp := self._typecheck(smt)) is not None:
+                self.expect_type(tp, VoidType(), smt)
+
+    @_node_typechecker(AstDeclNode)
+    def _typecheck_decl(self, n: AstDeclNode):
+        if not n.value:  # Nothing to check
+            return
+        expect = self._resolve_scope(n.scope).declared[n.ident.id].tp_info
+        self.expect_type(self._typecheck(n.value), expect, n)
+
+    @_node_typechecker(AstRepeat)
+    def _typecheck_repeat(self, n: AstRepeat):
+        # For now, we don't differentiate between number/string (as sc doesn't)
+        self.expect_type(self._typecheck(n.count), ValType(), n.count)
+        self._typecheck_block(n.body)
+
+    @_node_typechecker(AstIf)
+    def _typecheck_if(self, n: AstIf):
+        self.expect_type(self._typecheck(n.cond), BoolType(), n.cond)
+        self._typecheck_block(n.if_body)
+        if n.else_body is not None:
+            self._typecheck_block(n.else_body)
+
+    @_node_typechecker(AstWhile)
+    def _typecheck_while(self, n: AstWhile):
+        self.expect_type(self._typecheck(n.cond), BoolType(), n.cond)
+        self._typecheck_block(n.body)
+
+    @_node_typechecker(AstAssign)
+    def _typecheck_assign(self, n: AstAssign):  # super tempted to call this _typecheck_ass
+        if isinstance(n.target, AstIdent):
+            target_tp = self._curr_scope.used[n.target.id].tp_info
+        elif isinstance(n.target, AstItem):  # ls[i] = v
+            target_tp = self._typecheck(n.target)  # Also checks that `ls` is a list
+        elif isinstance(n.target, AstAttribute):
+            raise self.err("Setting attributes is currently unsupported", n.target)
+        else:
+            assert 0, "Unknown simple-assignment type"
+        if target_tp == ListType():
+            raise self.err("Cannot assign directly to list", n)
+        self.expect_type(self._typecheck(n.source), target_tp, n)
+
+    @_node_typechecker(AstAugAssign)
+    def _typecheck_aug_assign(self, n: AstAugAssign):
+        # TODO: change this when desugaring is implemented
+        #  (for now only +=, only on variables)
+        if n.op != '+=':
+            raise self.err(f"The '{n.op}' operator is not implemented", n)
+        if not isinstance(n.target, AstIdent):
+            raise self.err(f"The '+=' operator is only implemented for variables", n)
+        target_tp = self._curr_scope.used[n.target.id].tp_info
+        if target_tp != ValType():
+            raise self.err(f"Cannot apply += to {target_tp}", n)
+        self.expect_type(self._typecheck(n.source), ValType(), n.source)
+
+    @_node_typechecker(AstDefine)
+    def _typecheck_define(self, n: AstDefine):
+        # Don't really need to check much here - type is generated from the
+        # syntax so must be correct. Set _curr_scope and check body
+        func_info = self._curr_scope.declared[n.ident.id]
+        assert isinstance(func_info, FuncInfo)
+        old_scope = self._curr_scope
+        self._curr_scope = func_info.subscope
+        try:
+            self._typecheck_block(n.body)
+        finally:
+            self._curr_scope = old_scope
+
+    def _resolve_scope(self, scope_tp: VarDeclScope):
+        return self.top_scope if scope_tp == VarDeclScope.GLOBAL else self._curr_scope
+
+    def err(self, msg: str, loc: RegionUnionArgT):
+        return TypecheckError(msg, region_union(loc), self.src)
+
+    def expect_type(self, actual: TypeInfo, exp: TypeInfo, loc: RegionUnionArgT):
+        if exp != actual:
+            # TODO: maybe better type formatting
+            raise self.err(f"Expected type {exp}, got type {actual}", loc)