Fix various bug in the new typechecker module

This commit is a backport of many fixes that were done as part of
testing the new typecheck module that operates on the new AST
representation. This commit doesn't yet switch the default typechecker
to the new one, which requires to deeply change the `Cache` structures
and the LSP as well, which are still operating on the old
representation.

It also includes a bunch of new helpers for some AST data structures.

core/src/bytecode/ast/builder.rs

@@ -32,8 +32,6 @@ use super::{
     *,
 };
 
-use std::rc::Rc;
-
 type StaticPath = Vec<Ident>;
 
 /// Typestate style tag for `Field`s that are not yet completely specified
@@ -56,6 +54,12 @@ pub struct Field<'ast, State> {
     state: State,
     path: StaticPath,
     metadata: FieldMetadata<'ast>,
+    /// We need to store the documentation separately, because the metadata only accepts a string
+    /// that has been allocated in the AST allocator. We could require to thread the allocator for
+    /// calls to [Self::doc] or [Self::some_doc], but it's more ergonomic to keep the phases of
+    /// building a field and finalizing it separate. We thus store the documentation here
+    /// temporarily.
+    doc: Option<String>,
     /// As we might build contract element by element, we can't rely on
     /// `metadata.annotation.contracts`, which is a fixed array.
     ///
@@ -72,7 +76,7 @@ impl<'ast, A> Field<'ast, A> {
 
     /// Attach documentation metadata to the field, optionally
     pub fn some_doc(mut self, some_doc: Option<impl AsRef<str>>) -> Self {
-        self.metadata.doc = some_doc.map(|doc| Rc::from(doc.as_ref()));
+        self.doc = some_doc.map(|s| s.as_ref().to_string());
         self
     }
 
@@ -133,6 +137,7 @@ impl<'ast> Field<'ast, Incomplete> {
             state: Incomplete(),
             path: path.into_iter().map(|e| e.as_ref().into()).collect(),
             metadata: Default::default(),
+            doc: None,
             contracts: Vec::new(),
         }
     }
@@ -148,6 +153,7 @@ impl<'ast> Field<'ast, Incomplete> {
             state: Complete(None),
             path: self.path,
             metadata: self.metadata,
+            doc: self.doc,
             contracts: self.contracts,
         }
     }
@@ -158,6 +164,7 @@ impl<'ast> Field<'ast, Incomplete> {
             state: Complete(Some(value.into())),
             path: self.path,
             metadata: self.metadata,
+            doc: self.doc,
             contracts: self.contracts,
         }
     }
@@ -172,6 +179,7 @@ impl<'ast> Field<'ast, Complete<'ast>> {
             state: record,
             path: self.path,
             metadata: self.metadata,
+            doc: self.doc,
             contracts: self.contracts,
         };
         match value {
@@ -185,6 +193,7 @@ impl<'ast> Field<'ast, Record<'ast>> {
     /// Finalize the [`Field`] without setting a value
     pub fn no_value(mut self, alloc: &'ast AstAlloc) -> Record<'ast> {
         self.finalize_contracts(alloc);
+        self.metadata.doc = self.doc.map(|s| alloc.alloc_str(&s));
 
         self.state.field_defs.push(record::FieldDef {
             path: alloc.alloc_many(
@@ -202,6 +211,7 @@ impl<'ast> Field<'ast, Record<'ast>> {
     /// Finalize the [`Field`] by setting its a value
     pub fn value(mut self, alloc: &'ast AstAlloc, value: impl Into<Ast<'ast>>) -> Record<'ast> {
         self.finalize_contracts(alloc);
+        self.metadata.doc = self.doc.map(|s| alloc.alloc_str(&s));
 
         self.state.field_defs.push(record::FieldDef {
             path: alloc.alloc_many(
@@ -244,6 +254,7 @@ impl<'ast> Record<'ast> {
             state: self,
             path: vec![Ident::new(name)],
             metadata: Default::default(),
+            doc: None,
             contracts: Vec::new(),
         }
     }
@@ -270,6 +281,7 @@ impl<'ast> Record<'ast> {
             state: self,
             path: path.into_iter().map(|e| Ident::new(e)).collect(),
             metadata: Default::default(),
+            doc: None,
             contracts: Vec::new(),
         }
     }
@@ -511,7 +523,7 @@ mod tests {
                     (
                         "foo",
                         FieldMetadata {
-                            doc: Some(Rc::from("foo")),
+                            doc: Some(alloc.alloc_str("foo")),
                             ..Default::default()
                         },
                         None
@@ -520,7 +532,7 @@ mod tests {
                     (
                         "baz",
                         FieldMetadata {
-                            doc: Some(Rc::from("baz")),
+                            doc: Some(alloc.alloc_str("baz")),
                             ..Default::default()
                         },
                         None,
@@ -764,7 +776,7 @@ mod tests {
                 iter::once((
                     "foo",
                     FieldMetadata {
-                        doc: Some(Rc::from("foo?")),
+                        doc: Some(alloc.alloc_str("foo?")),
                         opt: true,
                         priority: MergePriority::Bottom,
                         not_exported: true,

core/src/bytecode/ast/compat.rs

@@ -186,7 +186,7 @@ impl<'ast> FromMainline<'ast, (LocIdent, term::record::Field)> for record::Field
 
 impl<'ast> FromMainline<'ast, term::record::FieldMetadata> for record::FieldMetadata<'ast> {
     fn from_mainline(alloc: &'ast AstAlloc, metadata: &term::record::FieldMetadata) -> Self {
-        let doc = metadata.doc.as_ref().map(|doc| rc::Rc::from(doc.as_str()));
+        let doc = metadata.doc.as_ref().map(|doc| alloc.alloc_str(doc));
 
         record::FieldMetadata {
             doc,
@@ -1019,6 +1019,15 @@ impl<'ast> FromAst<EnumRowsUnr<'ast>> for MainlineEnumRowsUnr {
     }
 }
 
+impl<'ast> FromAst<EnumRow<'ast>> for mline_type::EnumRow {
+    fn from_ast(erow: &EnumRow<'ast>) -> Self {
+        mline_type::EnumRow {
+            id: erow.id,
+            typ: erow.typ.as_ref().map(|ty| Box::new((*ty).to_mainline())),
+        }
+    }
+}
+
 impl<'ast> FromAst<RecordRowsUnr<'ast>> for MainlineRecordRowsUnr {
     fn from_ast(rrows: &RecordRowsUnr<'ast>) -> Self {
         rrows.clone().map(
@@ -1028,9 +1037,22 @@ impl<'ast> FromAst<RecordRowsUnr<'ast>> for MainlineRecordRowsUnr {
     }
 }
 
+impl<'ast> FromAst<RecordRow<'ast>> for mline_type::RecordRow {
+    fn from_ast(rrow: &RecordRow<'ast>) -> Self {
+        mline_type::RecordRowF {
+            id: rrow.id,
+            typ: Box::new(rrow.typ.to_mainline()),
+        }
+    }
+}
+
 impl<'ast> FromAst<Type<'ast>> for term::LabeledType {
     fn from_ast(typ: &Type<'ast>) -> Self {
         let typ: mline_type::Type = typ.to_mainline();
+        //TODO:remove
+        if typ.pos.into_opt().is_none() {
+            panic!("Expected a position to be set for the type {typ:?}");
+        }
         // We expect the new AST node to always have a position set. In fact we should
         // probably switch to `RawSpan` instead of `TermPos` everywhere; but let's do that
         // later

core/src/bytecode/ast/mod.rs

@@ -11,8 +11,7 @@
 
 use std::{
     ffi::{OsStr, OsString},
-    fmt::Debug,
-    iter, rc,
+    fmt, iter,
 };
 
 use pattern::Pattern;
@@ -165,7 +164,7 @@ pub struct LetBinding<'ast> {
 /// The metadata that can be attached to a let. It's a subset of [record::FieldMetadata].
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct LetMetadata<'ast> {
-    pub doc: Option<rc::Rc<str>>,
+    pub doc: Option<&'ast str>,
     pub annotation: Annotation<'ast>,
 }
 
@@ -233,6 +232,15 @@ impl Ast<'_> {
     }
 }
 
+impl Default for Ast<'_> {
+    fn default() -> Self {
+        Ast {
+            node: Node::Null,
+            pos: TermPos::None,
+        }
+    }
+}
+
 /// A branch of a match expression.
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct MatchBranch<'ast> {
@@ -263,18 +271,7 @@ pub struct Annotation<'ast> {
     pub contracts: &'ast [Type<'ast>],
 }
 
-impl<'ast> Annotation<'ast> {
-    /// Returns the main annotation, which is either the type annotation if any, or the first
-    /// contract annotation.
-    pub fn first<'a>(&'a self) -> Option<&'a Type<'ast>> {
-        self.typ.as_ref().or(self.contracts.iter().next())
-    }
-
-    /// Iterates over the annotations, starting by the type and followed by the contracts.
-    pub fn iter<'a>(&'a self) -> impl Iterator<Item = &'a Type<'ast>> {
-        self.typ.iter().chain(self.contracts.iter())
-    }
-
+impl Annotation<'_> {
     /// Returns a string representation of the contracts (without the static type annotation) as a
     /// comma-separated list.
     pub fn contracts_to_string(&self) -> Option<String> {
@@ -295,8 +292,8 @@ impl<'ast> Annotation<'ast> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
 /// Specifies where something should be imported from.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub enum Import<'ast> {
     Path {
         path: &'ast OsStr,
@@ -442,8 +439,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for Ast<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         state: &S,
     ) -> Option<U> {
         let child_state = match f(self, state) {
@@ -558,12 +555,12 @@ impl<'ast> TraverseAlloc<'ast, Type<'ast>> for Ast<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Type<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Type<'ast>, &S) -> TraverseControl<S, U>,
         state: &S,
     ) -> Option<U> {
         self.traverse_ref(
-            &mut |ast: &Ast<'ast>, state: &S| match &ast.node {
+            &mut |ast: &'ast Ast<'ast>, state: &S| match &ast.node {
                 Node::Type(typ) => typ.traverse_ref(f, state).into(),
                 _ => TraverseControl::Continue,
             },
@@ -592,8 +589,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for Annotation<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         scope: &S,
     ) -> Option<U> {
         self.typ
@@ -630,8 +627,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for LetBinding<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         scope: &S,
     ) -> Option<U> {
         self.metadata
@@ -666,8 +663,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for MatchBranch<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         scope: &S,
     ) -> Option<U> {
         self.pattern
@@ -697,6 +694,7 @@ impl Allocable for MatchBranch<'_> {}
 impl Allocable for Record<'_> {}
 impl Allocable for record::FieldPathElem<'_> {}
 impl Allocable for FieldDef<'_> {}
+impl Allocable for record::FieldMetadata<'_> {}
 
 impl Allocable for Pattern<'_> {}
 impl Allocable for EnumPattern<'_> {}
@@ -738,6 +736,11 @@ impl AstAlloc {
         }
     }
 
+    /// Return the current number of allocated bytes.
+    pub fn allocated_bytes(&self) -> usize {
+        self.generic_arena.allocated_bytes() + self.number_arena.len() + self.error_arena.len()
+    }
+
     /// Allocates an AST component in the arena.
     ///
     /// [Self] never guarantees that all destructors are going to be run when using such a generic
@@ -999,8 +1002,8 @@ impl AstAlloc {
 
 // Phony implementation of `Debug` so that we can still derive the trait for structure that holds
 // onto an allocator.
-impl Debug for AstAlloc {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl fmt::Debug for AstAlloc {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "AstAlloc")
     }
 }
@@ -1028,3 +1031,13 @@ where
 
     fn try_convert(alloc: &'ast AstAlloc, from: T) -> Result<Self, Self::Error>;
 }
+
+//TODO: get rid of this expensive implementation once we migrate pretty::*.
+impl fmt::Display for Ast<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use crate::term::RichTerm;
+        use compat::FromAst as _;
+
+        write!(f, "{}", RichTerm::from_ast(self))
+    }
+}

core/src/bytecode/ast/pattern/mod.rs

@@ -272,8 +272,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for Pattern<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         scope: &S,
     ) -> Option<U> {
         match &self.data {
@@ -324,8 +324,8 @@ impl<'ast> TraverseAlloc<'ast, Ast<'ast>> for FieldPattern<'ast> {
     }
 
     fn traverse_ref<S, U>(
-        &self,
-        f: &mut dyn FnMut(&Ast<'ast>, &S) -> TraverseControl<S, U>,
+        &'ast self,
+        f: &mut dyn FnMut(&'ast Ast<'ast>, &S) -> TraverseControl<S, U>,
         scope: &S,
     ) -> Option<U> {
         self.annotation