Fix function binding name resolution

fir/src/lib.rs

@@ -95,9 +95,9 @@
 // Does that make sense? Does that indicate that for all types we must first keep a Option<Ty> which is set to None?
 // Is this going to cause problems?
 
-use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::hash::Hash;
+use std::{collections::BTreeMap, ops::Index};
 
 mod checks;
 pub mod iter;
@@ -122,7 +122,7 @@ pub enum RefIdx {
 
 impl RefIdx {
     #[track_caller]
-    pub fn unwrap(&self) -> OriginIdx {
+    pub fn expect_resolved(&self) -> OriginIdx {
         match self {
             RefIdx::Resolved(idx) => *idx,
             RefIdx::Unresolved => unreachable!("unexpected `RefIdx::Unresolved` in `Fir`"),
@@ -223,6 +223,22 @@ pub enum Kind {
     Return(Option<RefIdx>),  // to any kind
 }
 
+impl<T> Index<&RefIdx> for Fir<T> {
+    type Output = Node<T>;
+
+    fn index(&self, index: &RefIdx) -> &Node<T> {
+        &self.nodes[&index.expect_resolved()]
+    }
+}
+
+impl<T> Index<&OriginIdx> for Fir<T> {
+    type Output = Node<T>;
+
+    fn index(&self, index: &OriginIdx) -> &Node<T> {
+        &self.nodes[index]
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Node<T = ()> {
     pub data: T,

flatten/src/lib.rs

@@ -666,7 +666,6 @@ impl<'ast> Ctx<'ast> {
 
     fn visit_type(
         self,
-
         ast: AstInfo<'ast>,
         generics: &[GenericArgument],
         fields: &[TypedValue],
@@ -690,7 +689,7 @@ impl<'ast> Ctx<'ast> {
     fn visit_var_declaration(
         self,
         ast: AstInfo<'ast>,
-        _mutable: &bool,
+        _mutable: bool,
         _to_declare: &Symbol,
         value: &'ast Ast,
     ) -> (Ctx<'ast>, RefIdx) {
@@ -746,6 +745,28 @@ impl<'ast> Ctx<'ast> {
         self.append(data, kind)
     }
 
+    fn handle_field_instantiation(self, instantiation: &'ast Ast) -> (Ctx<'ast>, RefIdx) {
+        let AstNode::VarAssign { value, .. } = &instantiation.node else {
+            // FIXME: Ugly?
+            unreachable!(
+                "invalid AST: non var-assign in field instantiation, in type instantiation"
+            )
+        };
+
+        let (ctx, value) = self.visit(value);
+
+        let data = FlattenData {
+            scope: ctx.scope,
+            ast: AstInfo::Node(instantiation),
+        };
+        let kind = Kind::Assignment {
+            to: RefIdx::Unresolved,
+            from: value,
+        };
+
+        ctx.append(data, kind)
+    }
+
     fn visit_type_instantiation(
         self,
         ast: AstInfo<'ast>,
@@ -756,7 +777,7 @@ impl<'ast> Ctx<'ast> {
         }: &'ast Call,
     ) -> (Ctx<'ast>, RefIdx) {
         let (ctx, generics) = self.visit_fold(generics.iter(), Ctx::handle_ty_node);
-        let (ctx, fields) = ctx.visit_fold(fields.iter(), Ctx::visit);
+        let (ctx, fields) = ctx.visit_fold(fields.iter(), Ctx::handle_field_instantiation);
 
         let data = FlattenData {
             scope: ctx.scope,
@@ -802,7 +823,7 @@ impl<'ast> Ctx<'ast> {
                 mutable,
                 to_declare,
                 value,
-            } => self.visit_var_declaration(node, mutable, to_declare, value),
+            } => self.visit_var_declaration(node, *mutable, to_declare, value),
             AstNode::VarAssign { to_assign, value } => {
                 self.visit_var_assign(node, to_assign, value)
             }
@@ -899,7 +920,7 @@ mod tests {
             Kind::Statements(stmts) => stmts,
             _ => unreachable!(),
         };
-        let ret_idx = stmts[0].unwrap();
+        let ret_idx = stmts[0].expect_resolved();
 
         assert!(matches!(
             fir.nodes.get(&ret_idx).unwrap().kind,

name_resolve/src/declarator.rs

@@ -1,60 +1,78 @@
-use fir::{Fallible, Fir, Node, RefIdx, Traversal};
+use fir::{Fallible, Fir, Node, OriginIdx, RefIdx, Traversal};
 use flatten::FlattenData;
 
-use crate::{NameResolutionError, NameResolveCtx};
+use crate::{NameResolutionError, NameResolveCtx, UniqueError};
 
-pub(crate) struct Declarator<'ctx>(pub(crate) &'ctx mut NameResolveCtx);
+enum DefinitionKind {
+    Function,
+    Type,
+    Binding,
+}
+
+pub(crate) struct Declarator<'ctx, 'enclosing>(pub(crate) &'ctx mut NameResolveCtx<'enclosing>);
+
+impl<'ctx, 'enclosing> Declarator<'ctx, 'enclosing> {
+    fn define(
+        &mut self,
+        kind: DefinitionKind,
+        node: &Node<FlattenData>,
+    ) -> Fallible<NameResolutionError> {
+        let (map, kind) = match kind {
+            DefinitionKind::Function => (&mut self.0.mappings.functions, "function"),
+            DefinitionKind::Type => (&mut self.0.mappings.types, "type"),
+            DefinitionKind::Binding => (&mut self.0.mappings.bindings, "binding"),
+        };
+
+        map.insert(
+            node.data.ast.symbol().unwrap().clone(),
+            node.origin,
+            self.0.enclosing_scope[node.origin],
+        )
+        .map_err(|existing| Declarator::unique_error(node, existing, kind))
+    }
+
+    fn unique_error(
+        node: &Node<FlattenData>,
+        existing: OriginIdx,
+        kind: &'static str,
+    ) -> NameResolutionError {
+        NameResolutionError::non_unique(node.data.ast.location(), UniqueError(existing, kind))
+    }
+}
+
+impl<'ast, 'ctx, 'enclosing> Traversal<FlattenData<'ast>, NameResolutionError>
+    for Declarator<'ctx, 'enclosing>
+{
+    // TODO: Can we factor these three functions?
 
-impl<'ast, 'ctx> Traversal<FlattenData<'ast>, NameResolutionError> for Declarator<'ctx> {
     fn traverse_function(
         &mut self,
-        _fir: &Fir<FlattenData>,
+        _: &Fir<FlattenData>,
         node: &Node<FlattenData>,
-        _generics: &[RefIdx],
-        _args: &[RefIdx],
-        _return_ty: &Option<RefIdx>,
-        _block: &Option<RefIdx>,
+        _: &[RefIdx],
+        _: &[RefIdx],
+        _: &Option<RefIdx>,
+        _: &Option<RefIdx>,
     ) -> Fallible<NameResolutionError> {
-        self.0
-            .mappings
-            .add_function(
-                node.data.ast.symbol().unwrap().clone(),
-                node.data.scope,
-                node.origin,
-            )
-            .map_err(|ue| NameResolutionError::non_unique(node.data.ast.location(), ue))
+        self.define(DefinitionKind::Function, node)
     }
 
     fn traverse_type(
         &mut self,
-        _fir: &Fir<FlattenData>,
+        _: &Fir<FlattenData>,
         node: &Node<FlattenData>,
         _: &[RefIdx],
         _: &[RefIdx],
     ) -> Fallible<NameResolutionError> {
-        self.0
-            .mappings
-            .add_type(
-                node.data.ast.symbol().unwrap().clone(),
-                node.data.scope,
-                node.origin,
-            )
-            .map_err(|ue| NameResolutionError::non_unique(node.data.ast.location(), ue))
+        self.define(DefinitionKind::Type, node)
     }
 
     fn traverse_binding(
         &mut self,
-        _fir: &Fir<FlattenData>,
+        _: &Fir<FlattenData>,
         node: &Node<FlattenData>,
-        _to: &RefIdx,
+        _: &RefIdx,
     ) -> Fallible<NameResolutionError> {
-        self.0
-            .mappings
-            .add_variable(
-                node.data.ast.symbol().unwrap().clone(),
-                node.data.scope,
-                node.origin,
-            )
-            .map_err(|ue| NameResolutionError::non_unique(node.data.ast.location(), ue))
+        self.define(DefinitionKind::Binding, node)
     }
 }