Reimplement optional type signature to force a type for a fn

src/lib/ast/tree.rs

@@ -51,7 +51,8 @@ impl Mod {
 
 #[derive(Debug, Clone)]
 pub enum TopLevel {
-    Prototype(Prototype),
+    Extern(Prototype),
+    FnSignature(Prototype),
     Function(FunctionDecl),
     Trait(Trait),
     Impl(Impl),
@@ -70,8 +71,12 @@ impl TopLevel {
         Self::Infix(op, pred)
     }
 
-    pub fn new_prototype(proto: Prototype) -> Self {
-        Self::Prototype(proto)
+    pub fn new_extern(proto: Prototype) -> Self {
+        Self::Extern(proto)
+    }
+
+    pub fn new_fn_signature(proto: Prototype) -> Self {
+        Self::FnSignature(proto)
     }
 
     pub fn new_use(u: Use) -> Self {

src/lib/ast/visit.rs

@@ -72,7 +72,8 @@ pub fn walk_mod<'a, V: Visitor<'a>>(visitor: &mut V, _mod: &'a Mod) {
 
 pub fn walk_top_level<'a, V: Visitor<'a>>(visitor: &mut V, top_level: &'a TopLevel) {
     match &top_level {
-        TopLevel::Prototype(p) => visitor.visit_prototype(p),
+        TopLevel::Extern(p) => visitor.visit_prototype(p),
+        TopLevel::FnSignature(p) => visitor.visit_prototype(p),
         TopLevel::Use(u) => visitor.visit_use(u),
         TopLevel::Trait(t) => visitor.visit_trait(t),
         TopLevel::Impl(i) => visitor.visit_impl(i),

src/lib/ast/visit_mut.rs

@@ -72,7 +72,8 @@ pub fn walk_mod<'a, V: VisitorMut<'a>>(visitor: &mut V, _mod: &'a mut Mod) {
 
 pub fn walk_top_level<'a, V: VisitorMut<'a>>(visitor: &mut V, top_level: &'a mut TopLevel) {
     match top_level {
-        TopLevel::Prototype(p) => visitor.visit_prototype(p),
+        TopLevel::Extern(p) => visitor.visit_prototype(p),
+        TopLevel::FnSignature(p) => visitor.visit_prototype(p),
         TopLevel::Use(u) => visitor.visit_use(u),
         TopLevel::Trait(t) => visitor.visit_trait(t),
         TopLevel::Impl(i) => visitor.visit_impl(i),

src/lib/ast_lowering/ast_lowering_context.rs

@@ -4,37 +4,43 @@ use crate::{
     ast::{tree::*, NodeId},
     hir::{self, Arena, FnBodyId, HirId},
     infer::Envs,
+    resolver::ResolutionMap,
     ty::*,
 };
 
 use super::{hir_map::HirMap, InfixDesugar};
 
 pub struct AstLoweringContext {
     hir_map: HirMap,
+    resolutions: ResolutionMap<NodeId>,
     top_levels: Vec<hir::TopLevel>,
     bodies: BTreeMap<FnBodyId, hir::FnBody>,
     operators_list: HashMap<String, u8>,
     traits: HashMap<Type, hir::Trait>,
     trait_methods: BTreeMap<String, HashMap<FuncType, hir::FunctionDecl>>,
     struct_methods: BTreeMap<HirId, HashMap<FuncType, hir::FunctionDecl>>,
     structs: HashMap<String, hir::StructDecl>,
+    signatures: ResolutionMap<NodeId>, // FnHirId -> SignatureHirId
 }
 
 impl AstLoweringContext {
     pub fn new(operators_list: HashMap<String, u8>) -> Self {
         Self {
             hir_map: HirMap::new(),
+            resolutions: ResolutionMap::new(),
             top_levels: Vec::new(),
             bodies: BTreeMap::new(),
             traits: HashMap::new(),
             trait_methods: BTreeMap::new(),
             struct_methods: BTreeMap::new(),
             structs: HashMap::new(),
+            signatures: ResolutionMap::new(),
             operators_list,
         }
     }
 
     pub fn lower_root(&mut self, root: &Root) -> hir::Root {
+        self.resolutions = root.resolutions.clone();
         self.lower_mod(&root.r#mod);
 
         let mut hir = hir::Root {
@@ -51,6 +57,11 @@ impl AstLoweringContext {
             spans: root.spans.clone(),
             structs: self.structs.clone(),
             trait_solver: root.trait_solver.clone(),
+            signatures: self
+                .signatures
+                .lower_resolution_map(&self.hir_map)
+                .inner()
+                .clone(),
         };
 
         hir.arena = hir::collect_arena(&hir);
@@ -67,13 +78,27 @@ impl AstLoweringContext {
 
     pub fn lower_top_level(&mut self, top_level: &TopLevel) {
         match &top_level {
-            TopLevel::Prototype(p) => {
+            TopLevel::Extern(p) => {
                 let top_level = hir::TopLevel {
-                    kind: hir::TopLevelKind::Prototype(self.lower_prototype(p)),
+                    kind: hir::TopLevelKind::Extern(self.lower_prototype(p)),
                 };
 
                 self.top_levels.push(top_level);
             }
+            TopLevel::FnSignature(p) => {
+                // TODO: gather existing signatures like this to have a pre-typecheck
+                let fn_node_id = self.resolutions.get(&p.node_id).unwrap();
+
+                let proto = self.lower_prototype(p);
+
+                let top_level = hir::TopLevel {
+                    kind: hir::TopLevelKind::Signature(proto),
+                };
+
+                self.top_levels.push(top_level);
+
+                self.signatures.insert(fn_node_id, p.node_id);
+            }
             TopLevel::Function(f) => {
                 let top_level = hir::TopLevel {
                     kind: hir::TopLevelKind::Function(self.lower_function_decl(f)),

src/lib/codegen/codegen_context.rs

@@ -165,7 +165,8 @@ impl<'a> CodegenContext<'a> {
     pub fn lower_hir(&mut self, root: &'a Root, builder: &'a Builder) -> Result<(), ()> {
         for item in &root.top_levels {
             match &item.kind {
-                TopLevelKind::Prototype(p) => self.lower_prototype(p, builder)?,
+                TopLevelKind::Extern(p) => self.lower_prototype(p, builder)?,
+                TopLevelKind::Signature(_p) => (),
                 TopLevelKind::Function(f) => self.lower_function_decl(f, builder)?,
             }
         }
@@ -646,9 +647,8 @@ impl<'a> CodegenContext<'a> {
 
         let callable_value = match self.hir.get_top_level(f_id) {
             Some(top) => CallableValue::try_from(match &top.kind {
-                TopLevelKind::Prototype(p) => {
-                    self.module.get_function(&p.name.to_string()).unwrap()
-                }
+                TopLevelKind::Extern(p) => self.module.get_function(&p.name.to_string()).unwrap(),
+                TopLevelKind::Signature(_p) => unimplemented!(),
                 TopLevelKind::Function(f) => {
                     self.module.get_function(&f.get_name().to_string()).unwrap()
                 }

src/lib/diagnostics/diagnostic.rs

@@ -84,6 +84,19 @@ impl Diagnostic {
         Self::new(span, DiagnosticKind::DuplicatedOperator)
     }
 
+    pub fn new_orphane_signature(span: Span, name: String) -> Self {
+        Self::new(span, DiagnosticKind::OrphaneSignature(name))
+    }
+
+    pub fn new_signature_mismatch(
+        span: Span,
+        name: String,
+        got: FuncType,
+        expected: FuncType,
+    ) -> Self {
+        Self::new(span, DiagnosticKind::SignatureMismatch(name, got, expected))
+    }
+
     pub fn new_no_main() -> Self {
         Self::new(Span::new_placeholder(), DiagnosticKind::NoMain)
     }
@@ -126,6 +139,8 @@ pub enum DiagnosticKind {
     CodegenError(HirId, String),
     IsNotAPropertyOf(Type, Span),
     OutOfBounds(u64, u64),
+    OrphaneSignature(String),
+    SignatureMismatch(String, FuncType, FuncType),
     NoMain,
     NoError, //TODO: remove that
 }
@@ -270,6 +285,20 @@ impl DiagnosticKind {
                         .with_message(format!("{}", self))
                         .with_color(color),
                 ),
+            DiagnosticKind::OrphaneSignature(name) => builder
+                .with_message(format!("Orpheline signature: {}", name))
+                .with_label(
+                    Label::new((span.file_path.to_str().unwrap(), span.start..span.end))
+                        .with_message(format!("{}", self))
+                        .with_color(color),
+                ),
+            DiagnosticKind::SignatureMismatch(name, _got, _expected) => builder
+                .with_message(format!("Signature mismatch: {}", name,))
+                .with_label(
+                    Label::new((span.file_path.to_str().unwrap(), span.start..span.end))
+                        .with_message(format!("{}", self))
+                        .with_color(color),
+                ),
             DiagnosticKind::NoMain => builder
                 .with_message("No main function".to_string())
                 .with_label(
@@ -342,6 +371,16 @@ impl Display for DiagnosticKind {
             DiagnosticKind::NotAFunction => "NotAFunction".to_string(),
             DiagnosticKind::UnusedParameter => "UnusedParameter".to_string(),
             DiagnosticKind::UnusedFunction => "UnusedFunction".to_string(),
+            DiagnosticKind::OrphaneSignature(_name) => "OrphelineSignature".to_string(),
+            DiagnosticKind::SignatureMismatch(_name, got, expected) => {
+                use colored::*;
+                format!(
+                    "Expected {}\n{:<9}But got  {}",
+                    format!("{:?}", expected).blue(),
+                    "",
+                    format!("{:?}", got).red(),
+                )
+            }
             DiagnosticKind::NoMain => "NoMain".to_string(),
             DiagnosticKind::NoError => "NoError".to_string(),
             DiagnosticKind::IsNotAPropertyOf(t, _span2) => {

src/lib/hir/tree.rs

@@ -10,7 +10,7 @@ use crate::{
     ty::{FuncType, StructType, Type},
 };
 
-use super::{arena::Arena, hir_printer, HasHirId};
+use super::{arena::Arena, hir_printer, HasHirId, HirNode};
 
 #[derive(Debug, Default)]
 pub struct Root {
@@ -27,6 +27,7 @@ pub struct Root {
     pub bodies: BTreeMap<FnBodyId, FnBody>,
     pub spans: HashMap<NodeId, Span>,
     pub structs: HashMap<String, StructDecl>,
+    pub signatures: HashMap<HirId, HirId>, // FnHirId -> SignatureHirId
 }
 
 impl Root {
@@ -80,6 +81,20 @@ impl Root {
             })
     }
 
+    pub fn get_function_signature_by_name(&self, name: &str) -> Option<FuncType> {
+        let fn_decl = self.get_function_by_name(name)?;
+
+        let sig_hir_id = self.signatures.get(&fn_decl.hir_id)?;
+
+        let arena_node = self.arena.get(sig_hir_id)?;
+
+        if let HirNode::Prototype(p) = arena_node {
+            Some(p.signature.clone())
+        } else {
+            None
+        }
+    }
+
     #[allow(dead_code)]
     pub fn get_function_by_mangled_name(&self, name: &str) -> Option<FunctionDecl> {
         self.top_levels
@@ -210,16 +225,18 @@ pub struct TopLevel {
 impl TopLevel {
     pub fn get_terminal_hir_id(&self) -> HirId {
         match &self.kind {
-            TopLevelKind::Prototype(p) => p.hir_id.clone(),
+            TopLevelKind::Extern(p) => p.hir_id.clone(),
             TopLevelKind::Function(f) => f.hir_id.clone(),
+            TopLevelKind::Signature(s) => s.hir_id.clone(),
         }
     }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum TopLevelKind {
     Function(FunctionDecl),
-    Prototype(Prototype),
+    Extern(Prototype),
+    Signature(Prototype),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

src/lib/hir/visit.rs

@@ -80,7 +80,8 @@ pub fn walk_root<'a, V: Visitor<'a>>(visitor: &mut V, root: &'a Root) {
 
 pub fn walk_top_level<'a, V: Visitor<'a>>(visitor: &mut V, top_level: &'a TopLevel) {
     match &top_level.kind {
-        TopLevelKind::Prototype(p) => visitor.visit_prototype(p),
+        TopLevelKind::Extern(p) => visitor.visit_prototype(p),
+        TopLevelKind::Signature(p) => visitor.visit_prototype(p),
         TopLevelKind::Function(f) => visitor.visit_function_decl(f),
     };
 }

src/lib/hir/visit_mut.rs

@@ -92,7 +92,8 @@ pub fn walk_struct_ctor<'a, V: VisitorMut<'a>>(visitor: &mut V, s: &'a mut Struc
 
 pub fn walk_top_level<'a, V: VisitorMut<'a>>(visitor: &mut V, top_level: &'a mut TopLevel) {
     match &mut top_level.kind {
-        TopLevelKind::Prototype(p) => visitor.visit_prototype(p),
+        TopLevelKind::Extern(p) => visitor.visit_prototype(p),
+        TopLevelKind::Signature(p) => visitor.visit_prototype(p),
         TopLevelKind::Function(f) => visitor.visit_function_decl(f),
     };
 }