From 3a26bb19249c89d16dd0b105c147c120c795200a Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Wed, 22 Jan 2025 12:49:33 -0800
Subject: [PATCH] fix: Analyze projection application

---
 Pantograph/Delate.lean | 38 +++++++++++++++++++-------------------
 Test/Delate.lean       | 25 ++++++++++++++++---------
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/Pantograph/Delate.lean b/Pantograph/Delate.lean
index a404fba..18eaaf1 100644
--- a/Pantograph/Delate.lean
+++ b/Pantograph/Delate.lean
@@ -12,27 +12,27 @@ open Lean
 
 namespace Pantograph
 
-structure ProjectionApplication where
-  projector: Name
-  numParams: Nat
-  inner: Expr
+inductive Projection where
+  -- Normal field case
+  | field (projector : Name) (numParams : Nat) (struct : Expr)
+  -- Singular inductive case
+  | singular (recursor : Name) (numParams : Nat) (struct : Expr)
 
-/-- Converts a `.proj` expression to a function application if possible. Not all
-such expressions are convertible. -/
-@[export pantograph_expr_proj_to_app]
-def exprProjToApp (env: Environment) (e: Expr): Option ProjectionApplication := do
-  let (typeName, idx, inner) := match e with
-    | .proj typeName idx inner => (typeName, idx, inner)
+/-- Converts a `.proj` expression to a form suitable for exporting/transpilation -/
+@[export pantograph_analyze_projection]
+def analyzeProjection (env: Environment) (e: Expr): Projection :=
+  let (typeName, idx, struct) := match e with
+    | .proj typeName idx struct => (typeName, idx, struct)
     | _ => panic! "Argument must be proj"
-  let _ ← getStructureInfo? env typeName
-  let ctor := getStructureCtor env typeName
-  let fieldName := getStructureFields env typeName |>.get! idx
-  let projector := getProjFnForField? env typeName fieldName |>.get!
-  return {
-    projector,
-    numParams := ctor.numParams,
-    inner,
-  }
+  if (getStructureInfo? env typeName).isSome then
+    let ctor := getStructureCtor env typeName
+    let fieldName := getStructureFields env typeName |>.get! idx
+    let projector := getProjFnForField? env typeName fieldName |>.get!
+    .field projector ctor.numParams struct
+  else
+    let recursor := mkRecOnName typeName
+    let ctor := getStructureCtor env typeName
+    .singular recursor ctor.numParams struct
 
 def _root_.Lean.Name.isAuxLemma (n : Lean.Name) : Bool := n matches .num (.str _ "_auxLemma") _
 
diff --git a/Test/Delate.lean b/Test/Delate.lean
index 7b449bf..aa5ea7c 100644
--- a/Test/Delate.lean
+++ b/Test/Delate.lean
@@ -96,17 +96,23 @@ def test_instance (env: Environment): IO LSpec.TestSeq :=
     let _expr := (← runTermElabMInMeta <| elabTerm s) |>.toOption |>.get!
     return LSpec.TestSeq.done
 
-def test_projection (env: Environment) : IO LSpec.TestSeq:= runTest do
-  let prod := .app (.bvar 1) (.bvar 0)
-  let expr := .proj `Prod 1 prod
-  let .some { projector, numParams, inner }:= exprProjToApp env expr |
-    fail "`Prod should have projection function"
+def test_projection_prod (env: Environment) : IO LSpec.TestSeq:= runTest do
+  let struct := .app (.bvar 1) (.bvar 0)
+  let expr := .proj `Prod 1 struct
+  let .field projector numParams struct' := analyzeProjection env expr |
+    fail "`Prod has fields"
   checkEq "projector" projector `Prod.snd
   checkEq "numParams" numParams 2
-  checkTrue "inner" $ inner.equal prod
+  checkTrue "struct" $ struct.equal struct'
 
-  let expr := .proj `Exists 1 prod
-  checkTrue "Exists" (exprProjToApp env expr).isNone
+def test_projection_exists (env: Environment) : IO LSpec.TestSeq:= runTest do
+  let struct := .app (.bvar 1) (.bvar 0)
+  let expr := .proj `Exists 1 struct
+  let .singular recursor numParams struct' := analyzeProjection env expr |
+    fail "`Exists has no projectors"
+  checkEq "recursor" recursor `Exists.recOn
+  checkEq "numParams" numParams 2
+  checkTrue "struct" $ struct.equal struct'
 
 def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
   [
@@ -116,7 +122,8 @@ def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
     ("Sexp from elaborated expr", test_sexp_of_elab env),
     ("Sexp from expr", test_sexp_of_expr env),
     ("Instance", test_instance env),
-    ("Projection", test_projection env),
+    ("Projection Prod", test_projection_prod env),
+    ("Projection Exists", test_projection_exists env),
   ]
 
 end Pantograph.Test.Delate