Unify the various typechecking functions

src/solvers/bm_based.rs

@@ -104,29 +104,11 @@ mod convert {
     }
 }
 
-pub fn trace_with_formality<'a>(conf: Conf, req: &TypingRequest<'a>) -> String {
-    let stmt = req.to_bm_based();
-    let mut r = Reduction::from_stmt(conf, stmt);
-    let mut out = String::new();
-    loop {
-        if r.last {
-            if !r.is_err() {
-                r.apply_dbm();
-                let _ = write!(&mut out, "{}", r);
-            }
-            break;
-        } else {
-            let _ = write!(&mut out, "{}", r);
-            r.step();
-        }
-    }
-    out
-}
-
-pub fn typecheck_with_formality<'a>(
+pub fn run_formality<'a>(
     a: &'a Arenas<'a>,
     conf: Conf,
     req: &TypingRequest<'a>,
+    mut callback: impl FnMut(&Reduction),
 ) -> TypingResult<'a> {
     if req.pat.contains_abstract() {
         return TypingResult::TypeError(TypeError::OverlyGeneral(DeepeningRequest::Pattern));
@@ -135,14 +117,41 @@ pub fn typecheck_with_formality<'a>(
     }
 
     let stmt = req.to_bm_based();
-    let r = Reduction::from_stmt(conf, stmt);
-    match r.to_type() {
-        Ok((ident, ty)) => {
+    let mut r = Reduction::from_stmt(conf, stmt);
+    while !r.last {
+        callback(&r);
+        r.step();
+    }
+
+    match r.node_step.error {
+        Some(e) => TypingResult::TypeError(TypeError::External(e)),
+        None => {
+            let (ident, ty) = r.as_type();
             let ty = Type::from_bm_based(a, &ty);
             let name = a.alloc_str(&ident.name);
             let bindings = BindingAssignments::new([(name, ty)]);
-            Success(bindings)
+            let out = Success(bindings);
+            // Just for printing, feels risky to do before `as_type`.
+            r.apply_dbm();
+            callback(&r);
+            out
         }
-        Err(e) => TypingResult::TypeError(TypeError::External(e)),
     }
 }
+
+pub fn trace_with_formality<'a>(conf: Conf, req: &TypingRequest<'a>) -> String {
+    let a = &Arenas::default();
+    let mut out = String::new();
+    run_formality(a, conf, req, |r| {
+        let _ = write!(&mut out, "{}", r);
+    });
+    out
+}
+
+pub fn typecheck_with_formality<'a>(
+    a: &'a Arenas<'a>,
+    conf: Conf,
+    req: &TypingRequest<'a>,
+) -> TypingResult<'a> {
+    run_formality(a, conf, req, |_| {})
+}

src/solvers/mod.rs

@@ -33,10 +33,14 @@ pub enum TypingResult<'a> {
 }
 
 impl RuleSet {
-    pub fn typecheck<'a>(&self, a: &'a Arenas<'a>, req: &TypingRequest<'a>) -> TypingResult<'a> {
+    pub fn typecheck<'a>(
+        &self,
+        arenas: &'a Arenas<'a>,
+        req: &TypingRequest<'a>,
+    ) -> TypingResult<'a> {
         match *self {
-            RuleSet::TypeBased(options) => typecheck_with_this_crate(a, options, req),
-            RuleSet::BindingModeBased(conf) => typecheck_with_formality(a, conf, req),
+            RuleSet::TypeBased(options) => typecheck_with_this_crate(arenas, options, req),
+            RuleSet::BindingModeBased(conf) => typecheck_with_formality(arenas, conf, req),
         }
     }
 

src/solvers/ty_based.rs

@@ -84,6 +84,53 @@ impl<'a> TypingSolver<'a> {
     }
 }
 
+pub enum SolverTraceEvent<'a, 'b> {
+    Start,
+    Step(Rule),
+    CantStep(&'b CantStep<'a>),
+}
+
+/// Run the solver on this request and return the result of typechecking.
+pub fn run_solver<'a>(
+    ctx: TypingCtx<'a>,
+    request: &TypingRequest<'a>,
+    mut callback: impl FnMut(&TypingSolver<'a>, SolverTraceEvent<'a, '_>),
+) -> TypingResult<'a> {
+    let mut solver = TypingSolver::new(*request);
+    callback(&solver, SolverTraceEvent::Start);
+    let e = loop {
+        match solver.step(ctx) {
+            Ok(rule) => callback(&solver, SolverTraceEvent::Step(rule)),
+            Err(e) => {
+                callback(&solver, SolverTraceEvent::CantStep(&e));
+                break e;
+            }
+        }
+    };
+
+    match e {
+        CantStep::Done => {
+            let bindings = BindingAssignments::new(solver.done_predicates.iter().map(|pred| {
+                let ty = *pred.expr.ty;
+                let Pattern::Binding(_, _, name) = pred.pat else {
+                    unreachable!()
+                };
+                (*name, ty)
+            }));
+            let borrow_check: Result<(), _> = solver
+                .done_predicates
+                .iter()
+                .map(|pred| pred.expr.simplify(ctx).borrow_check())
+                .collect();
+            match borrow_check {
+                Ok(()) => TypingResult::Success(bindings),
+                Err(err) => TypingResult::BorrowError(bindings, err),
+            }
+        }
+        CantStep::NoApplicableRule(_, err) => TypingResult::TypeError(err),
+    }
+}
+
 /// Run the solver on this request and returns the trace as a string.
 pub fn trace_solver<'a>(
     request: TypingRequest<'a>,
@@ -92,65 +139,35 @@ pub fn trace_solver<'a>(
 ) -> String {
     let arenas = &Arenas::default();
     let ctx = TypingCtx { arenas, options };
-    let mut solver = TypingSolver::new(request);
     let mut trace = String::new();
-    let _ = write!(&mut trace, "{}\n", solver.display_state(style));
-    loop {
-        match solver.step(ctx) {
-            Ok(rule) => {
-                let line = format!("// Applying rule `{}`", rule.display(options));
-                let _ = write!(&mut trace, "{}\n", line.comment());
-                let _ = write!(&mut trace, "{}\n", solver.display_state(style));
+    run_solver(ctx, &request, |solver, event| match event {
+        SolverTraceEvent::Start => {
+            let _ = write!(&mut trace, "{}\n", solver.display_state(style));
+        }
+        SolverTraceEvent::Step(rule) => {
+            let line = format!("// Applying rule `{}`", rule.display(options));
+            let _ = write!(&mut trace, "{}\n", line.comment());
+            let _ = write!(&mut trace, "{}\n", solver.display_state(style));
+        }
+        SolverTraceEvent::CantStep(e) => match e {
+            CantStep::Done => {
+                let _ = write!(&mut trace, "\n// Final bindings (simplified):\n");
+                let _ = write!(&mut trace, "{}\n", solver.display_final_state(ctx, style));
             }
-            Err(e) => {
-                match e {
-                    CantStep::Done => {
-                        let _ = write!(&mut trace, "\n// Final bindings (simplified):\n");
-                        let _ = write!(&mut trace, "{}\n", solver.display_final_state(ctx, style));
-                    }
-                    CantStep::NoApplicableRule(pred, err) => {
-                        let line = format!("// Type error for `{}`: {err:?}", pred.display(style));
-                        let _ = write!(&mut trace, "{}\n", line.red());
-                    }
-                }
-                break;
+            CantStep::NoApplicableRule(pred, err) => {
+                let line = format!("// Type error for `{}`: {err:?}", pred.display(style));
+                let _ = write!(&mut trace, "{}\n", line.red());
             }
-        }
-    }
+        },
+    });
     trace
 }
 
 pub fn typecheck_with_this_crate<'a>(
-    a: &'a Arenas<'a>,
+    arenas: &'a Arenas<'a>,
     options: RuleOptions,
     req: &TypingRequest<'a>,
 ) -> TypingResult<'a> {
-    let ctx = TypingCtx { arenas: a, options };
-    let mut solver = TypingSolver::new(*req);
-    // TODO: abstract over the repeated stepping of the solver
-    let e = loop {
-        match solver.step(ctx) {
-            Ok(_) => {}
-            Err(e) => break e,
-        }
-    };
-    match e {
-        CantStep::Done => {
-            assert_eq!(solver.done_predicates.len(), 1);
-            let pred = solver.done_predicates[0];
-            let ty = *pred.expr.ty;
-            let Pattern::Binding(_, _, name) = pred.pat else {
-                unreachable!()
-            };
-            let bindings = BindingAssignments::new([(*name, ty)]);
-            match pred.expr.simplify(ctx).borrow_check() {
-                // This error isn't handled by `match-ergo-formality` so we ignore it.
-                Ok(()) | Err(BorrowCheckError::CantCopyNestedRefMut) => {
-                    TypingResult::Success(bindings)
-                }
-                Err(err) => TypingResult::BorrowError(bindings, err),
-            }
-        }
-        CantStep::NoApplicableRule(_, err) => TypingResult::TypeError(err),
-    }
+    let ctx = TypingCtx { arenas, options };
+    run_solver(ctx, req, |_, _| {})
 }