Elaborate letrec terms

Author: Kmeakin

fathom/src/core.rs

@@ -143,13 +143,8 @@ pub enum Term<'arena> {
     /// Annotated expressions.
     Ann(Span, &'arena Term<'arena>, &'arena Term<'arena>),
     /// Let expressions.
-    Let(
-        Span,
-        Option<StringId>,
-        &'arena Term<'arena>,
-        &'arena Term<'arena>,
-        &'arena Term<'arena>,
-    ),
+    Let(Span, &'arena LetDef<'arena>, &'arena Term<'arena>),
+    Letrec(Span, &'arena [LetDef<'arena>], &'arena Term<'arena>),
 
     /// The type of types.
     Universe(Span),
@@ -204,6 +199,13 @@ pub enum Term<'arena> {
     ),
 }
 
+#[derive(Debug, Clone)]
+pub struct LetDef<'arena> {
+    pub name: Option<StringId>,
+    pub r#type: Term<'arena>,
+    pub expr: Term<'arena>,
+}
+
 impl<'arena> Term<'arena> {
     /// Get the source span of the term.
     pub fn span(&self) -> Span {
@@ -213,7 +215,8 @@ impl<'arena> Term<'arena> {
             | Term::MetaVar(span, _)
             | Term::InsertedMeta(span, _, _)
             | Term::Ann(span, _, _)
-            | Term::Let(span, _, _, _, _)
+            | Term::Let(span, ..)
+            | Term::Letrec(span, ..)
             | Term::Universe(span)
             | Term::FunType(span, ..)
             | Term::FunLit(span, ..)
@@ -243,11 +246,23 @@ impl<'arena> Term<'arena> {
             | Term::ConstLit(_, _) => false,
 
             Term::Ann(_, expr, r#type) => expr.binds_local(var) || r#type.binds_local(var),
-            Term::Let(_, _, def_type, def_expr, body_expr) => {
-                def_type.binds_local(var)
-                    || def_expr.binds_local(var)
+            Term::Let(_, def, body_expr) => {
+                def.r#type.binds_local(var)
+                    || def.expr.binds_local(var)
                     || body_expr.binds_local(var.prev())
             }
+            Term::Letrec(_, defs, body_expr) => {
+                for _def in defs.iter() {
+                    var = var.prev();
+                }
+
+                if (defs.iter()).any(|def| def.r#type.binds_local(var) || def.expr.binds_local(var))
+                {
+                    return true;
+                }
+
+                body_expr.binds_local(var)
+            }
             Term::FunType(.., param_type, body_type) => {
                 param_type.binds_local(var) || body_type.binds_local(var.prev())
             }

fathom/src/core/pretty.rs

@@ -30,7 +30,7 @@ use std::cell::RefCell;
 
 use pretty::RcDoc;
 
-use crate::core::{Item, Module, Plicity, Term};
+use crate::core::{Item, LetDef, Module, Plicity, Term};
 use crate::source::{StringId, StringInterner};
 use crate::surface::lexer::is_keyword;
 
@@ -150,24 +150,33 @@ impl<'interner, 'arena> Context<'interner> {
                     self.term_prec(Prec::Top, r#type),
                 ]),
             ),
-            Term::Let(_, def_pattern, def_type, def_expr, body_expr) => self.paren(
+            Term::Let(_, def, body_expr) => self.paren(
                 prec > Prec::Let,
                 RcDoc::concat([
                     RcDoc::concat([
                         RcDoc::text("let"),
                         RcDoc::space(),
-                        self.ann_pattern(Prec::Top, *def_pattern, def_type),
-                        RcDoc::space(),
-                        RcDoc::text("="),
-                        RcDoc::softline(),
-                        self.term_prec(Prec::Let, def_expr),
+                        self.let_def(def),
                         RcDoc::text(";"),
                     ])
                     .group(),
                     RcDoc::line(),
                     self.term_prec(Prec::Let, body_expr),
                 ]),
             ),
+            Term::Letrec(_, defs, body_expr) => self.paren(
+                prec > Prec::Let,
+                self.sequence(
+                    RcDoc::concat([RcDoc::text("let")]),
+                    defs.iter().map(|def| self.let_def(def)),
+                    RcDoc::text(","),
+                    RcDoc::concat([
+                        RcDoc::text(";"),
+                        RcDoc::line(),
+                        self.term_prec(Prec::Let, body_expr),
+                    ]),
+                ),
+            ),
             Term::Universe(_) => RcDoc::text("Type"),
             Term::FunType(_, plicity, param_name, param_type, body_type) => self.paren(
                 prec > Prec::Fun,
@@ -341,6 +350,16 @@ impl<'interner, 'arena> Context<'interner> {
         }
     }
 
+    fn let_def(&'arena self, def: &LetDef<'arena>) -> RcDoc {
+        RcDoc::concat([
+            self.ann_pattern(Prec::Top, def.name, &def.r#type),
+            RcDoc::space(),
+            RcDoc::text("="),
+            RcDoc::softline(),
+            self.term_prec(Prec::Let, &def.expr),
+        ])
+    }
+
     fn format_field(&'arena self, label: StringId, format: &Term<'arena>) -> RcDoc {
         RcDoc::concat([
             self.ident(label),

fathom/src/core/semantics.rs

@@ -7,8 +7,8 @@ use std::sync::Arc;
 use scoped_arena::Scope;
 
 use crate::alloc::SliceVec;
-use crate::core::{prim, Const, LocalInfo, Plicity, Prim, Term};
-use crate::env::{EnvLen, Index, Level, SharedEnv, SliceEnv};
+use crate::core::*;
+use crate::env::{self, EnvLen, Index, Level, SharedEnv, SliceEnv};
 use crate::source::{Span, Spanned, StringId};
 
 /// Atomically reference counted values. We use reference counting to increase
@@ -53,6 +53,8 @@ pub enum Value<'arena> {
 }
 
 impl<'arena> Value<'arena> {
+    pub const ERROR: Self = Self::Stuck(Head::Prim(Prim::ReportedError), Vec::new());
+
     pub fn prim(prim: Prim, params: impl IntoIterator<Item = ArcValue<'arena>>) -> Value<'arena> {
         let params = params
             .into_iter()
@@ -299,13 +301,29 @@ impl<'arena, 'env> EvalEnv<'arena, 'env> {
                 self.apply_local_infos(head_expr, local_infos)
             }
             Term::Ann(span, expr, _) => Spanned::merge(*span, self.eval(expr)),
-            Term::Let(span, _, _, def_expr, body_expr) => {
-                let def_expr = self.eval(def_expr);
+            Term::Let(span, def, body_expr) => {
+                let def_expr = self.eval(&def.expr);
                 self.local_exprs.push(def_expr);
                 let body_expr = self.eval(body_expr);
                 self.local_exprs.pop();
                 Spanned::merge(*span, body_expr)
             }
+            Term::Letrec(span, defs, body_expr) => {
+                let initial_len = self.local_exprs.len();
+
+                for _def in defs.iter() {
+                    // TODO: lazy evaluation?
+                    (self.local_exprs).push(Spanned::empty(Arc::new(Value::ERROR)));
+                }
+                for (level, def) in Iterator::zip(env::levels(), defs.iter()) {
+                    let def_expr = self.eval(&def.expr);
+                    self.local_exprs.set_level(level, def_expr);
+                }
+
+                let body_expr = self.eval(body_expr);
+                self.local_exprs.truncate(initial_len);
+                Spanned::merge(*span, body_expr)
+            }
 
             Term::Universe(span) => Spanned::new(*span, Arc::new(Value::Universe)),
 
@@ -879,13 +897,35 @@ impl<'arena, 'env> EvalEnv<'arena, 'env> {
                 scope.to_scope(self.unfold_metas(scope, expr)),
                 scope.to_scope(self.unfold_metas(scope, r#type)),
             ),
-            Term::Let(span, def_name, def_type, def_expr, body_expr) => Term::Let(
+            Term::Let(span, def, body_expr) => Term::Let(
                 *span,
-                *def_name,
-                scope.to_scope(self.unfold_metas(scope, def_type)),
-                scope.to_scope(self.unfold_metas(scope, def_expr)),
+                scope.to_scope(LetDef {
+                    name: def.name,
+                    r#type: self.unfold_metas(scope, &def.r#type),
+                    expr: self.unfold_metas(scope, &def.expr),
+                }),
                 self.unfold_bound_metas(scope, body_expr),
             ),
+            Term::Letrec(span, defs, body_expr) => {
+                let initial_len = self.local_exprs.len();
+
+                for _def in defs.iter() {
+                    let var = Arc::new(Value::local_var(self.local_exprs.len().next_level()));
+                    self.local_exprs.push(Spanned::empty(var));
+                }
+
+                let defs = scope.to_scope_from_iter(defs.iter().map(|def| {
+                    let name = def.name;
+                    let r#type = self.unfold_metas(scope, &def.r#type);
+                    let expr = self.unfold_metas(scope, &def.expr);
+                    LetDef { name, r#type, expr }
+                }));
+
+                let body_expr = self.unfold_metas(scope, body_expr);
+                self.local_exprs.truncate(initial_len);
+
+                Term::Letrec(*span, defs, scope.to_scope(body_expr))
+            }
 
             Term::Universe(span) => Term::Universe(*span),
 

fathom/src/env.rs

@@ -337,6 +337,10 @@ impl<Entry> SharedEnv<Entry> {
         self.get_level(self.len().index_to_level(index)?)
     }
 
+    pub fn set_level(&mut self, level: Level, entry: Entry) {
+        self.entries.set_mut(level.0 as usize, entry);
+    }
+
     /// Push an entry onto the environment.
     pub fn push(&mut self, entry: Entry) {
         assert!(self.entries.len() < usize::from(u16::MAX));

fathom/src/surface.rs

@@ -179,6 +179,15 @@ impl<Range: Clone> Pattern<Range> {
     }
 }
 
+impl<Range> Pattern<Range> {
+    pub fn name(&self) -> Option<StringId> {
+        match self {
+            Pattern::Name(_, name) => Some(*name),
+            _ => None,
+        }
+    }
+}
+
 /// Surface terms.
 #[derive(Debug, Clone)]
 pub enum Term<'arena, Range> {

fathom/src/surface/distillation.rs

@@ -380,10 +380,10 @@ impl<'interner, 'arena, 'env> Context<'interner, 'arena, 'env> {
                     Term::Ann((), self.scope.to_scope(expr), self.scope.to_scope(r#type)),
                 )
             }
-            (core::Term::Let(_, name, r#type, expr, body), _) => {
-                let r#type = self.term_prec(mode, Prec::Top, r#type);
-                let expr = self.term_prec(mode, Prec::Let, expr);
-                let name = self.freshen_name(*name, body);
+            (core::Term::Let(_, def, body), _) => {
+                let r#type = self.term_prec(mode, Prec::Top, &def.r#type);
+                let expr = self.term_prec(mode, Prec::Let, &def.expr);
+                let name = self.freshen_name(def.name, body);
                 let name = self.push_local(name);
                 let pattern = name_to_pattern(name);
                 let body = self.term_prec(mode, Prec::Top, body);
@@ -402,6 +402,30 @@ impl<'interner, 'arena, 'env> Context<'interner, 'arena, 'env> {
                     ),
                 )
             }
+            (core::Term::Letrec(_, defs, body_expr), _) => {
+                let initial_len = self.local_len();
+
+                // TODO: freshen names
+                let def_names: Vec<_> = defs.iter().map(|def| self.push_local(def.name)).collect();
+
+                let defs = self.scope.to_scope_from_iter(
+                    Iterator::zip(def_names.iter(), defs.iter()).map(|(name, def)| {
+                        let r#type = self.check(&def.r#type);
+                        let expr = self.check(&def.expr);
+
+                        LetDef {
+                            pattern: name_to_pattern(*name),
+                            r#type: Some(r#type),
+                            expr,
+                        }
+                    }),
+                );
+
+                let body_expr = self.check(body_expr);
+                self.truncate_local(initial_len);
+
+                Term::Letrec((), defs, self.scope.to_scope(body_expr))
+            }
             (core::Term::Universe(_), _) => Term::Universe(()),
             (core::Term::FunType(..), _) => {
                 let initial_local_len = self.local_len();