More type checking -- expressions

Compiler/Ast.fir

@@ -79,6 +79,14 @@ type TypeDecl(
 )
 
 
+impl Tokens[TypeDecl]:
+    firstToken(self: TypeDecl) TokenIdx:
+        self.prim_.unwrapOr(self.name.token)
+
+    lastToken(self: TypeDecl) TokenIdx:
+        panic("TOOD")
+
+
 ## Constructors of a type declaration.
 #[derive(ToDoc)]
 type TypeDeclRhs:
@@ -98,6 +106,14 @@ type ConDecl(
 )
 
 
+impl Tokens[ConDecl]:
+    firstToken(self: ConDecl) TokenIdx:
+        self.name.token
+
+    lastToken(self: ConDecl) TokenIdx:
+        self._lastToken
+
+
 #[derive(ToDoc)]
 type ConFields:
     Empty
@@ -724,7 +740,7 @@ type Constructor(
     userTyArgs: Option[TyArgs],
 
     ## Inferred type arguments of the constructor's type. Filled in by the type checker.
-    tyArgs: Vec[Type],
+    tyArgs: Vec[Ty],
 
     # In variant constructors we'll see a '~' before `ty`, so we can't use `ty` as the first token.
     _firstToken: TokenIdx,
@@ -873,7 +889,7 @@ type VarExpr(
     userTyArgs: Option[TyArgs],
 
     ## Inferred type arguments of the variable. Filled in by the type checker.
-    tyArgs: Vec[Type],
+    tyArgs: Vec[Ty],
 )
 
 

Compiler/Defs.fir

@@ -53,6 +53,12 @@ VarId.resolve(self, def: VarDefIdx):
     self._resolved = Option.Some(def)
 
 
+VarId.def(self) VarDefIdx:
+    self._resolved.unwrapOrElse(
+        \(): panic("Term id (\``self.name`\`) is not resolved yet"),
+    )
+
+
 #[derive(ToDoc)]
 type AssocVarId(
     ## Index of the associated member identifier's token in its module.
@@ -161,30 +167,12 @@ type TyDefIdx(
 
     ## Name of the type in `_mod`.
     _name: Str,
-
-    ## Kind of the definition: type or trait.
-    _kind: TyDefKind,
 )
 
 
-#[derive(ToDoc)]
-type TyDefKind:
-    Type
-    Trait
-
-
 impl Eq[TyDefIdx]:
     __eq(self: TyDefIdx, other: TyDefIdx) Bool:
-        self._mod == other._mod
-            and self._name == other._name
-            and self._kind == other._kind
-
-
-impl Eq[TyDefKind]:
-    __eq(self: TyDefKind, other: TyDefKind) Bool:
-        match self:
-            TyDefKind.Type: other is TyDefKind.Type
-            TyDefKind.Trait: other is TyDefKind.Trait
+        self._mod == other._mod and self._name == other._name
 
 
 # ------------------------------------------------------------------------------
@@ -317,21 +305,10 @@ impl ToDoc[TyDefIdx]:
                 + Doc.nested(4, Doc.break_(1) + self._name.toDoc()),
         )
         args += Doc.char(',') + Doc.break_(1)
-        args += Doc.grouped(
-            Doc.str("_kind =")
-                + Doc.nested(4, Doc.break_(1) + self._kind.toDoc()),
-        )
         args = args.nest(4).group() + Doc.break_(0) + Doc.char(')')
         Doc.grouped(Doc.str("TyDefIdx") + Doc.char('(') + args)
 
 
-impl ToDoc[TyDefKind]:
-    toDoc(self: TyDefKind) Doc:
-        match self:
-            TyDefKind.Type: Doc.str("TyDefKind.Type")
-            TyDefKind.Trait: Doc.str("TyDefKind.Trait")
-
-
 # ------------------------------------------------------------------------------
 # ToStr implementations. These are used when generating error messages, so they
 # should be readable by the users and should not expose implementation details.

Compiler/Module.fir

@@ -85,7 +85,7 @@ type Module(
     # modules.
 
     ## Term environment of the module, used for name resolving.
-    _termEnv: HashMap[Str, VarDefIdx],
+    _termEnv: HashMap[Str, TopVarDefIdx],
 
     ## Associated item environment of the module, used for name resolving.
     _assocTermEnv: HashMap[Str, HashMap[Str, AssocVarDefIdx]],
@@ -124,12 +124,44 @@ Module.loc[Tokens[t]](self, t: t) Loc:
 
 Module.tyCon(self, tyIdx: TyDefIdx) TyCon:
     assert(self._idx == tyIdx._mod)
-    let tcEnv = self._tcEnv.unwrap()
-    let item = tcEnv._cons.get(tyIdx._name).unwrap()
+
+    let tcEnv = self._tcEnv.unwrapOrElse(
+        \(): panic("Module.tyCon called before preparing type environments"),
+    )
+
+    # `TyDefIdx`s are assigned by the name resolver after finding the
+    # definition, so it's a bug if a `TyDefIdx` is not in the environment.
+    let item = tcEnv._cons.get(tyIdx._name).unwrapOrElse(
+        \():
+            panic(
+                "BUG: Module.tyCon TyDefIdx is not in the module type checking environment",
+            ),
+    )
+
     assert(item.idx == tyIdx)
     item.info
 
 
+Module.scheme(self, varIdx: TopVarDefIdx) Scheme:
+    assert(self._idx == varIdx._mod)
+
+    let tcEnv = self._tcEnv.unwrapOrElse(
+        \(): panic("Module.scheme called before preparing type environments"),
+    )
+
+    # `TopVarDefIdx`s are assigned by the name resolver after finding the
+    # definition, so it's a bug if a `TopVarDefIdx` is not in the environment.
+    let item = tcEnv._topSchemes.get(varIdx._name).unwrapOrElse(
+        \():
+            panic(
+                "BUG: Module.scheme TopVarDefIdx is not in the module type checking environment",
+            ),
+    )
+
+    assert(item.idx == varIdx)
+    item.info
+
+
 # ------------------------------------------------------------------------------
 # Generated ToDoc implementations
 

Compiler/NameResolver.fir

@@ -15,7 +15,14 @@ resolveNames(pgm: Program):
     _prepRenamingEnvs(pgm)
 
     for module: Module in pgm._modules.iter():
-        let termEnv = ScopeMap.fromMap(module._termEnv)
+        let termEnv = ScopeMap.fromMap(
+            HashMap.fromIter(
+                module._termEnv.iter().map(
+                    \(entry: HashMapEntry[Str, TopVarDefIdx]):
+                        (key = entry.key, value = VarDefIdx.Top(entry.value)),
+                ),
+            ),
+        )
 
         let visitor = _NameResolverVisitor(
             _moduleFilePath = module._filePath,
@@ -220,14 +227,14 @@ _makeSccDepGraph(pgm: Program, sccs: Vec[HashSet[ModuleIdx]]) Vec[HashSet[U32]]:
 
 
 type _SccEnv(
-    _termEnv: HashMap[Str, VarDefIdx],
+    _termEnv: HashMap[Str, TopVarDefIdx],
     _assocTermEnv: HashMap[Str, HashMap[Str, AssocVarDefIdx]],
     _tyEnv: HashMap[Str, TyDefIdx],
 )
 
 
 _SccEnv.merge(self, other: _SccEnv):
-    for term: HashMapEntry[Str, VarDefIdx] in other._termEnv.iter():
+    for term: HashMapEntry[Str, TopVarDefIdx] in other._termEnv.iter():
         let old = self._termEnv.insert(term.key, term.value)
         if old is Option.Some(oldImport) and oldImport != term.value:
             panic("Variable `term.key` defined multiple times in SCC")
@@ -264,7 +271,7 @@ _makeSccEnv(
     if sccEnvs.get(sccIdx) is Option.Some(env):
         return env
 
-    let termEnv: HashMap[Str, VarDefIdx] = HashMap.withCapacity(50)
+    let termEnv: HashMap[Str, TopVarDefIdx] = HashMap.withCapacity(50)
     let assocTermEnv: HashMap[Str, HashMap[Str, AssocVarDefIdx]] = HashMap
         .withCapacity(25)
     let typeEnv: HashMap[Str, TyDefIdx] = HashMap.withCapacity(25)
@@ -276,11 +283,7 @@ _makeSccEnv(
         let tokens = mod._tokens
 
         for ty: HashMapEntry[Str, TypeDecl] in mod._tyItems.iter():
-            let tyDefIdx = TyDefIdx(
-                _mod = modIdx,
-                _name = ty.key,
-                _kind = TyDefKind.Type,
-            )
+            let tyDefIdx = TyDefIdx(_mod = modIdx, _name = ty.key,)
 
             let old = typeEnv.insert(ty.key, tyDefIdx)
             if old is Option.Some(_):
@@ -289,21 +292,16 @@ _makeSccEnv(
                 )
 
         for trait_: HashMapEntry[Str, TraitDecl] in mod._traitItems.iter():
-            let traitDefIdx = TyDefIdx(
-                _mod = modIdx,
-                _name = trait_.key,
-                _kind = TyDefKind.Trait,
-            )
+            let traitDefIdx = TyDefIdx(_mod = modIdx, _name = trait_.key,)
             let old = typeEnv.insert(trait_.key, traitDefIdx)
             if old is Option.Some(_):
                 panic(
                     "Type `trait_.key` defined multiple times in a recursive import group",
                 )
 
         for fun: HashMapEntry[Str, FunDecl] in mod._funItems.iter():
-            let varDefIdx = VarDefIdx.Top(
-                TopVarDefIdx(_mod = modIdx, _name = fun.key),
-            )
+            let varDefIdx = TopVarDefIdx(_mod = modIdx, _name = fun.key)
+
             let old = termEnv.insert(fun.key, varDefIdx)
             if old is Option.Some(_):
                 panic(

Compiler/Program.fir

@@ -27,11 +27,20 @@ Program.new() Program:
     Program(_modules = Vec.empty(), _depGraph = Option.None,)
 
 
+Program.module(self, modIdx: ModuleIdx) Module:
+    self._modules.get(modIdx._idx)
+
+
 Program.tyCon(self, tyIdx: TyDefIdx) TyCon:
     let mod = self._modules.get(tyIdx._mod._idx)
     mod.tyCon(tyIdx)
 
 
+Program.scheme(self, varIdx: TopVarDefIdx) Scheme:
+    let mod = self._modules.get(varIdx._mod._idx)
+    mod.scheme(varIdx)
+
+
 ## Load a module, or return it from the cache if it's already loaded.
 ##
 ## Follows imports in the module to transitively import all referred modules.