summaryrefslogtreecommitdiff
path: root/backends/lean/Base/Diverge/Elab.lean
blob: 0088fd16b684e66c3ab1ff5f64cec280449601e6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
import Lean
import Lean.Meta.Tactic.Simp
import Init.Data.List.Basic
import Mathlib.Tactic.RunCmd
import Base.Utils
import Base.Diverge.Base
import Base.Diverge.ElabBase

namespace Diverge

/- Automating the generation of the encoding and the proofs so as to use nice
   syntactic sugar. -/

syntax (name := divergentDef)
  declModifiers "divergent" "def" declId ppIndent(optDeclSig) declVal : command

open Lean Elab Term Meta Primitives Lean.Meta
open Utils

/- The following was copied from the `wfRecursion` function. -/

open WF in

def mkProd (x y : Expr) : MetaM Expr :=
  mkAppM ``Prod.mk #[x, y]

def mkInOutTy (x y : Expr) : MetaM Expr :=
  mkAppM ``FixI.mk_in_out_ty #[x, y]

-- Return the `a` in `Return a`
def getResultTy (ty : Expr) : MetaM Expr :=
  ty.withApp fun f args => do
  if ¬ f.isConstOf ``Result  args.size  1 then
    throwError "Invalid argument to getResultTy: {ty}"
  else
    pure (args.get! 0)

/- Deconstruct a sigma type.

   For instance, deconstructs `(a : Type) × List a` into
   `Type` and `λ a => List a`.
 -/
def getSigmaTypes (ty : Expr) : MetaM (Expr × Expr) := do
  ty.withApp fun f args => do
  if ¬ f.isConstOf ``Sigma  args.size  2 then
    throwError "Invalid argument to getSigmaTypes: {ty}"
  else
    pure (args.get! 0, args.get! 1)

/- Generate a Sigma type from a list of *variables* (all the expressions
   must be variables).

   Example:
   - xl = [(a:Type), (ls:List a), (i:Int)]

   Generates:
   `(a:Type) × (ls:List a) × (i:Int)`

 -/
def mkSigmasType (xl : List Expr) : MetaM Expr :=
  match xl with
  | [] => do
    trace[Diverge.def.sigmas] "mkSigmasOfTypes: []"
    pure (Expr.const ``PUnit.unit [])
  | [x] => do
    trace[Diverge.def.sigmas] "mkSigmasOfTypes: [{x}]"
    let ty  Lean.Meta.inferType x
    pure ty
  | x :: xl => do
    trace[Diverge.def.sigmas] "mkSigmasOfTypes: [{x}::{xl}]"
    let alpha  Lean.Meta.inferType x
    let sty  mkSigmasType xl
    trace[Diverge.def.sigmas] "mkSigmasOfTypes: [{x}::{xl}]: alpha={alpha}, sty={sty}"
    let beta  mkLambdaFVars #[x] sty
    trace[Diverge.def.sigmas] "mkSigmasOfTypes: ({alpha}) ({beta})"
    mkAppOptM ``Sigma #[some alpha, some beta]

/- Apply a lambda expression to some arguments, simplifying the lambdas -/
def applyLambdaToArgs (e : Expr) (xs : Array Expr) : MetaM Expr := do
  lambdaTelescopeN e xs.size fun vars body =>
  -- Create the substitution
  let s : HashMap FVarId Expr := HashMap.ofList (List.zip (vars.toList.map Expr.fvarId!) xs.toList)
  -- Substitute in the body
  pure (body.replace fun e =>
    match e with
    | Expr.fvar fvarId => match s.find? fvarId with
      | none   => e
      | some v => v
    | _ => none)

/- Group a list of expressions into a dependent tuple.

   Example:
   xl = [`a : Type`, `ls : List a`]
   returns:
   `⟨ (a:Type), (ls: List a) ⟩`

   We need the type argument because as the elements in the tuple are
   "concrete", we can't in all generality figure out the type of the tuple.

   Example:
   `⟨ True, 3 ⟩ : (x : Bool) × (if x then Int else Unit)`
 -/
def mkSigmasVal (ty : Expr) (xl : List Expr) : MetaM Expr :=
  match xl with
  | [] => do
    trace[Diverge.def.sigmas] "mkSigmasVal: []"
    pure (Expr.const ``PUnit.unit [])
  | [x] => do
    trace[Diverge.def.sigmas] "mkSigmasVal: [{x}]"
    pure x
  | fst :: xl => do
    trace[Diverge.def.sigmas] "mkSigmasVal: [{fst}::{xl}]"
    -- Deconstruct the type
    let (alpha, beta)  getSigmaTypes ty
    -- Compute the "second" field
    -- Specialize beta for fst
    let nty  applyLambdaToArgs beta #[fst]
    -- Recursive call
    let snd  mkSigmasVal nty xl
    -- Put everything together
    trace[Diverge.def.sigmas] "mkSigmasVal:\n{alpha}\n{beta}\n{fst}\n{snd}"
    mkAppOptM ``Sigma.mk #[some alpha, some beta, some fst, some snd]

def mkAnonymous (s : String) (i : Nat) : Name :=
  .num (.str .anonymous s) i

/- Given a list of values `[x0:ty0, ..., xn:ty1]`, where every `xi` might use the previous
   `xj` (j < i) and a value `out` which uses `x0`, ..., `xn`, generate the following
   expression:
   ```
   fun x:((x0:ty0) × ... × (xn:tyn) => -- **Dependent** tuple
   match x with
   | (x0, ..., xn) => out
   ```

   The `index` parameter is used for naming purposes: we use it to numerotate the
   bound variables that we introduce.

   We use this function to currify functions (the function bodies given to the
   fixed-point operator must be unary functions).

   Example:
   ========
   - xl = `[a:Type, ls:List a, i:Int]`
   - out = `a`
   - index = 0

   generates (getting rid of most of the syntactic sugar):
   ```
   λ scrut0 => match scrut0 with
   | Sigma.mk x scrut1 =>
     match scrut1 with
     | Sigma.mk ls i =>
       a
   ```
-/
partial def mkSigmasMatch (xl : List Expr) (out : Expr) (index : Nat := 0) : MetaM Expr :=
  match xl with
  | [] => do
    -- This would be unexpected
    throwError "mkSigmasMatch: empyt list of input parameters"
  | [x] => do
    -- In the example given for the explanations: this is the inner match case
    trace[Diverge.def.sigmas] "mkSigmasMatch: [{x}]"
    mkLambdaFVars #[x] out
  | fst :: xl => do
    -- In the example given for the explanations: this is the outer match case
    -- Remark: for the naming purposes, we use the same convention as for the
    -- fields and parameters in `Sigma.casesOn` and `Sigma.mk` (looking at
    -- those definitions might help)
    --
    -- We want to build the match expression:
    -- ```
    -- λ scrut =>
    -- match scrut with
    -- | Sigma.mk x ...  -- the hole is given by a recursive call on the tail
    -- ```
    trace[Diverge.def.sigmas] "mkSigmasMatch: [{fst}::{xl}]"
    let alpha  Lean.Meta.inferType fst
    let snd_ty  mkSigmasType xl
    let beta  mkLambdaFVars #[fst] snd_ty
    let snd  mkSigmasMatch xl out (index + 1)
    let mk  mkLambdaFVars #[fst] snd
    -- Introduce the "scrut" variable
    let scrut_ty  mkSigmasType (fst :: xl)
    withLocalDeclD (mkAnonymous "scrut" index) scrut_ty fun scrut => do
    trace[Diverge.def.sigmas] "mkSigmasMatch: scrut: ({scrut}) : ({← inferType scrut})"
    -- TODO: make the computation of the motive more efficient
    let motive  do
      let out_ty  inferType out
      match out_ty  with
      | .sort _ | .lit _ | .const .. =>
        -- The type of the motive doesn't depend on the scrutinee
        mkLambdaFVars #[scrut] out_ty
      | _ =>
        -- The type of the motive *may* depend on the scrutinee
        -- TODO: make this more efficient (we could change the output type of
        -- mkSigmasMatch
        mkSigmasMatch (fst :: xl) out_ty
    -- The final expression: putting everything together
    trace[Diverge.def.sigmas] "mkSigmasMatch:\n  ({alpha})\n  ({beta})\n  ({motive})\n  ({scrut})\n  ({mk})"
    let sm  mkAppOptM ``Sigma.casesOn #[some alpha, some beta, some motive, some scrut, some mk]
    -- Abstracting the "scrut" variable
    let sm  mkLambdaFVars #[scrut] sm
    trace[Diverge.def.sigmas] "mkSigmasMatch: sm: {sm}"
    pure sm

/- Small tests for list_nth: give a model of what `mkSigmasMatch` should generate -/
private def list_nth_out_ty_inner (a :Type) (scrut1: @Sigma (List a) (fun (_ls : List a) => Int)) :=
  @Sigma.casesOn (List a)
                 (fun (_ls : List a) => Int)
                 (fun (_scrut1:@Sigma (List a) (fun (_ls : List a) => Int)) => Type)
                 scrut1
                 (fun (_ls : List a) (_i : Int) => Primitives.Result a)

private def list_nth_out_ty_outer (scrut0 : @Sigma (Type) (fun (a:Type) =>
                      @Sigma (List a) (fun (_ls : List a) => Int))) :=
  @Sigma.casesOn (Type)
                 (fun (a:Type) => @Sigma (List a) (fun (_ls : List a) => Int))
                 (fun (_scrut0:@Sigma (Type) (fun (a:Type) => @Sigma (List a) (fun (_ls : List a) => Int))) => Type)
                 scrut0
                 (fun (a : Type) (scrut1: @Sigma (List a) (fun (_ls : List a) => Int)) =>
                  list_nth_out_ty_inner a scrut1)
/- -/

-- Return the expression: `Fin n`
-- TODO: use more
def mkFin (n : Nat) : Expr :=
  mkAppN (.const ``Fin []) #[.lit (.natVal n)]

-- Return the expression: `i : Fin n`
def mkFinVal (n i : Nat) : MetaM Expr := do
  let n_lit : Expr := .lit (.natVal (n - 1))
  let i_lit : Expr := .lit (.natVal i)
  -- We could use `trySynthInstance`, but as we know the instance that we are
  -- going to use, we can save the lookup
  let ofNat  mkAppOptM ``Fin.instOfNatFinHAddNatInstHAddInstAddNatOfNat #[n_lit, i_lit]
  mkAppOptM ``OfNat.ofNat #[none, none, ofNat]

/- Generate and declare as individual definitions the bodies for the individual funcions:
   - replace the recursive calls with calls to the continutation `k`
   - make those bodies take one single dependent tuple as input

   We name the declarations: "[original_name].body".
   We return the new declarations.
 -/
def mkDeclareUnaryBodies (grLvlParams : List Name) (kk_var : Expr)
  (inOutTys : Array (Expr × Expr)) (preDefs : Array PreDefinition) :
  MetaM (Array Expr) := do
  let grSize := preDefs.size

  -- Compute the map from name to (index × input type).
  -- Remark: the continuation has an indexed type; we use the index (a finite number of
  -- type `Fin`) to control which function we call at the recursive call site.
  let nameToInfo : HashMap Name (Nat × Expr) :=
    let bl := preDefs.mapIdx fun i d => (d.declName, (i.val, (inOutTys.get! i.val).fst))
    HashMap.ofList bl.toList

  trace[Diverge.def.genBody] "nameToId: {nameToInfo.toList}"

  -- Auxiliary function to explore the function bodies and replace the
  -- recursive calls
  let visit_e (i : Nat) (e : Expr) : MetaM Expr := do
    trace[Diverge.def.genBody] "visiting expression (dept: {i}): {e}"
    let ne  do
      match e with
      | .app .. => do
        e.withApp fun f args => do
          trace[Diverge.def.genBody] "this is an app: {f} {args}"
          -- Check if this is a recursive call
          if f.isConst then
            let name := f.constName!
            match nameToInfo.find? name with
            | none => pure e
            | some (id, in_ty) =>
              trace[Diverge.def.genBody] "this is a recursive call"
              -- This is a recursive call: replace it
              -- Compute the index
              let i  mkFinVal grSize id
              -- Put the arguments in one big dependent tuple
              let args  mkSigmasVal in_ty args.toList
              mkAppM' kk_var #[i, args]
          else
            -- Not a recursive call: do nothing
            pure e
       | .const name _ =>
         -- Sanity check: we eliminated all the recursive calls
         if (nameToInfo.find? name).isSome then
           throwError "mkUnaryBodies: a recursive call was not eliminated"
         else pure e
       | _ => pure e
    trace[Diverge.def.genBody] "done with expression (depth: {i}): {e}"
    pure ne

  -- Explore the bodies
  preDefs.mapM fun preDef => do
    -- Replace the recursive calls
    trace[Diverge.def.genBody] "About to replace recursive calls in {preDef.declName}"
    let body  mapVisit visit_e preDef.value
    trace[Diverge.def.genBody] "Body after replacement of the recursive calls: {body}"

    -- Currify the function by grouping the arguments into a dependent tuple
    -- (over which we match to retrieve the individual arguments).
    lambdaTelescope body fun args body => do
    let body  mkSigmasMatch args.toList body 0

    -- Add the declaration
    let value  mkLambdaFVars #[kk_var] body
    let name := preDef.declName.append "body"
    let levelParams := grLvlParams
    let decl := Declaration.defnDecl {
      name := name
      levelParams := levelParams
      type :=  inferType value -- TODO: change the type
      value := value
      hints := ReducibilityHints.regular (getMaxHeight ( getEnv) value + 1)
      safety := .safe
      all := [name]
    }
    addDecl decl
    trace[Diverge.def] "individual body of {preDef.declName}: {body}"
    -- Return the constant
    let body := Lean.mkConst name (levelParams.map .param)
    -- let body ← mkAppM' body #[kk_var]
    trace[Diverge.def] "individual body (after decl): {body}"
    pure body

-- Generate a unique function body from the bodies of the mutually recursive group,
-- and add it as a declaration in the context.
-- We return the list of bodies (of type `FixI.Funs ...`) and the mutually recursive body.
def mkDeclareMutRecBody (grName : Name) (grLvlParams : List Name)
  (kk_var i_var : Expr)
  (in_ty out_ty : Expr) (inOutTys : List (Expr × Expr))
  (bodies : Array Expr) : MetaM (Expr × Expr) := do
  -- Generate the body
  let grSize := bodies.size
  let finTypeExpr := mkFin grSize
  -- TODO: not very clean
  let inOutTyType  do
    let (x, y) := inOutTys.get! 0
    inferType ( mkInOutTy x y)
  let rec mkFuns (inOutTys : List (Expr × Expr)) (bl : List Expr) : MetaM Expr :=
    match inOutTys, bl with
    | [], [] =>
      mkAppOptM ``FixI.Funs.Nil #[finTypeExpr, in_ty, out_ty]
    | (ity, oty) :: inOutTys, b :: bl => do
      -- Retrieving ity and oty - this is not very clean
      let inOutTysExpr  mkListLit inOutTyType ( inOutTys.mapM (λ (x, y) => mkInOutTy x y))
      let fl  mkFuns inOutTys bl
      mkAppOptM ``FixI.Funs.Cons #[finTypeExpr, in_ty, out_ty, ity, oty, inOutTysExpr, b, fl]
    | _, _ => throwError "mkDeclareMutRecBody: `tys` and `bodies` don't have the same length"
  let bodyFuns  mkFuns inOutTys bodies.toList
  -- Wrap in `get_fun`
  let body  mkAppM ``FixI.get_fun #[bodyFuns, i_var, kk_var]
  -- Add the index `i` and the continuation `k` as a variables
  let body  mkLambdaFVars #[kk_var, i_var] body
  trace[Diverge.def] "mkDeclareMutRecBody: body: {body}"
  -- Add the declaration
  let name := grName.append "mut_rec_body"
  let levelParams := grLvlParams
  let decl := Declaration.defnDecl {
    name := name
    levelParams := levelParams
    type :=  inferType body
    value := body
    hints := ReducibilityHints.regular (getMaxHeight ( getEnv) body + 1)
    safety := .safe
    all := [name]
  }
  addDecl decl
  -- Return the bodies and the constant
  pure (bodyFuns, Lean.mkConst name (levelParams.map .param))

def isCasesExpr (e : Expr) : MetaM Bool := do
  let e := e.getAppFn
  if e.isConst then
    return isCasesOnRecursor ( getEnv) e.constName
  else return false

structure MatchInfo where
  matcherName       : Name
  matcherLevels     : Array Level
  params            : Array Expr
  motive            : Expr
  scruts            : Array Expr
  branchesNumParams : Array Nat
  branches          : Array Expr

instance : ToMessageData MatchInfo where
  -- This is not a very clean formatting, but we don't need more
  toMessageData := fun me => m!"\n- matcherName: {me.matcherName}\n- params: {me.params}\n- motive: {me.motive}\n- scruts: {me.scruts}\n- branchesNumParams: {me.branchesNumParams}\n- branches: {me.branches}"

-- Small helper: prove that an expression which doesn't use the continuation `kk`
-- is valid, and return the proof.
def proveNoKExprIsValid (k_var : Expr) (e : Expr) : MetaM Expr := do
  trace[Diverge.def.valid] "proveNoKExprIsValid: {e}"
  let eIsValid  mkAppM ``FixI.is_valid_p_same #[k_var, e]
  trace[Diverge.def.valid] "proveNoKExprIsValid: result:\n{eIsValid}:\n{← inferType eIsValid}"
  pure eIsValid

mutual

/- Prove that an expression is valid, and return the proof.

   More precisely, if `e` is an expression which potentially uses the continution
   `kk`, return an expression of type:
   ```
   is_valid_p k (λ kk => e)
   ```
 -/
partial def proveExprIsValid (k_var kk_var : Expr) (e : Expr) : MetaM Expr := do
  trace[Diverge.def.valid] "proveValid: {e}"
  match e with
  | .const _ _ => throwError "Unimplemented" -- Shouldn't get there?
  | .bvar _
  | .fvar _
  | .lit _
  | .mvar _
  | .sort _ => throwError "Unreachable"
  | .lam .. => throwError "Unimplemented" -- TODO
  | .forallE .. => throwError "Unreachable" -- Shouldn't get there
  | .letE .. => do
    -- Telescope all the let-bindings (remark: this also telescopes the lambdas)
    lambdaLetTelescope e fun xs body => do
    -- Note that we don't visit the bound values: there shouldn't be
    -- recursive calls, lambda expressions, etc. inside
    -- Prove that the body is valid
    let isValid  proveExprIsValid k_var kk_var body
    -- Add the let-bindings around.
    -- Rem.: the let-binding should be *inside* the `is_valid_p`, not outside,
    -- but because it reduces in the end it doesn't matter. More precisely:
    -- `P (let x := v in y)` and `let x := v in P y` reduce to the same expression.
    mkLambdaFVars xs isValid (usedLetOnly := false)
  | .mdata _ b => proveExprIsValid k_var kk_var b
  | .proj _ _ _ =>
    -- The projection shouldn't use the continuation
    proveNoKExprIsValid k_var e
  | .app .. =>
    e.withApp fun f args => do
      -- There are several cases: first, check if this is a match/if
      -- Check if the expression is a (dependent) if then else.
      -- We treat the if then else expressions differently from the other matches,
      -- and have dedicated theorems for them.
      let isIte := e.isIte
      if isIte || e.isDIte then do
        e.withApp fun f args => do
        trace[Diverge.def.valid] "ite/dite: {f}:\n{args}"
        if args.size  5 then
           throwError "Wrong number of parameters for {f}: {args}"
        let cond := args.get! 1
        let dec := args.get! 2
        -- Prove that the branches are valid
        let br0 := args.get! 3
        let br1 := args.get! 4
        let proveBranchValid (br : Expr) : MetaM Expr :=
          if isIte then proveExprIsValid k_var kk_var br
          else do
            -- There is a lambda
            lambdaOne br fun x br => do
            let brValid  proveExprIsValid k_var kk_var br
            mkLambdaFVars #[x] brValid
        let br0Valid  proveBranchValid br0
        let br1Valid  proveBranchValid br1
        let const := if isIte then ``FixI.is_valid_p_ite else ``FixI.is_valid_p_dite
        let eIsValid  mkAppOptM const #[none, none, none, none, some k_var, some cond, some dec, none, none, some br0Valid, some br1Valid]
        trace[Diverge.def.valid] "ite/dite: result:\n{eIsValid}:\n{← inferType eIsValid}"
        pure eIsValid
      -- Check if the expression is a match (this case is for when the elaborator
      -- introduces auxiliary definitions to hide the match behind syntactic
      -- sugar):
      else if let some me :=  matchMatcherApp? e then do
        trace[Diverge.def.valid]
          "matcherApp:
           - params: {me.params}
           - motive: {me.motive}
           - discrs: {me.discrs}
           - altNumParams: {me.altNumParams}
           - alts: {me.alts}
           - remaining: {me.remaining}"
        -- matchMatcherApp does all the work for us: we simply need to gather
        -- the information and call the auxiliary helper `proveMatchIsValid`
        if me.remaining.size  0 then
          throwError "MatcherApp: non empty remaining array: {me.remaining}"
        let me : MatchInfo := {
          matcherName := me.matcherName
          matcherLevels := me.matcherLevels
          params := me.params
          motive := me.motive
          scruts := me.discrs
          branchesNumParams := me.altNumParams
          branches := me.alts
        }
        proveMatchIsValid k_var kk_var me
      -- Check if the expression is a raw match (this case is for when the expression
      -- is a direct call to the primitive `casesOn` function, without syntactic sugar).
      -- We have to check this case because functions like `mkSigmasMatch`, which we
      -- use to currify function bodies, introduce such raw matches.
      else if  isCasesExpr f then do
        trace[Diverge.def.valid] "rawMatch: {e}"
        -- Deconstruct the match, and call the auxiliary helper `proveMatchIsValid`.
        --
        -- The casesOn definition is always of the following shape:
        -- - input parameters (implicit parameters)
        -- - motive (implicit), -- the motive gives the return type of the match
        -- - scrutinee (explicit)
        -- - branches (explicit).
        -- In particular, we notice that the scrutinee is the first *explicit*
        -- parameter - this is how we spot it.
        let matcherName := f.constName!
        let matcherLevels := f.constLevels!.toArray
        -- Find the first explicit parameter: this is the scrutinee
        forallTelescope ( inferType f) fun xs _ => do
        let rec findFirstExplicit (i : Nat) : MetaM Nat := do
          if i  xs.size then throwError "Unexpected: could not find an explicit parameter"
          else
            let x := xs.get! i
            let xFVarId := x.fvarId!
            let localDecl  xFVarId.getDecl
            match localDecl.binderInfo with
            | .default => pure i
            | _ => findFirstExplicit (i + 1)
        let scrutIdx  findFirstExplicit 0
        -- Split the arguments
        let params := args.extract 0 (scrutIdx - 1)
        let motive := args.get! (scrutIdx - 1)
        let scrut := args.get! scrutIdx
        let branches := args.extract (scrutIdx + 1) args.size
        -- Compute the number of parameters for the branches: for this we use
        -- the type of the uninstantiated casesOn constant (we can't just
        -- destruct the lambdas in the branch expressions because the result
        -- of a match might be a lambda expression).
        let branchesNumParams : Array Nat  do
          let env  getEnv
          let decl := env.constants.find! matcherName
          let ty := decl.type
          forallTelescope ty fun xs _ => do
          let xs := xs.extract (scrutIdx + 1) xs.size
          xs.mapM fun x => do
          let xty  inferType x
          forallTelescope xty fun ys _ => do
          pure ys.size
        let me : MatchInfo := {
          matcherName,
          matcherLevels,
          params,
          motive,
          scruts := #[scrut],
          branchesNumParams,
          branches,
        }
        proveMatchIsValid k_var kk_var me
      -- Check if this is a monadic let-binding
      else if f.isConstOf ``Bind.bind then do
        trace[Diverge.def.valid] "bind:\n{args}"
        -- We simply need to prove that the subexpressions are valid, and call
        -- the appropriate lemma.
        let x := args.get! 4
        let y := args.get! 5
        -- Prove that the subexpressions are valid
        let xValid  proveExprIsValid k_var kk_var x
        trace[Diverge.def.valid] "bind: xValid:\n{xValid}:\n{← inferType xValid}"
        let yValid  do
          -- This is a lambda expression
          lambdaOne y fun x y => do
          trace[Diverge.def.valid] "bind: y: {y}"
          let yValid  proveExprIsValid k_var kk_var y
          trace[Diverge.def.valid] "bind: yValid (no forall): {yValid}"
          trace[Diverge.def.valid] "bind: yValid: x: {x}"
          let yValid  mkLambdaFVars #[x] yValid
          trace[Diverge.def.valid] "bind: yValid (forall): {yValid}: {← inferType yValid}"
          pure yValid
        -- Put everything together
        trace[Diverge.def.valid] "bind:\n- xValid: {xValid}: {← inferType xValid}\n- yValid: {yValid}: {← inferType yValid}"
        mkAppM ``FixI.is_valid_p_bind #[xValid, yValid]
      -- Check if this is a recursive call, i.e., a call to the continuation `kk`
      else if f.isFVarOf kk_var.fvarId! then do
        trace[Diverge.def.valid] "rec: args: \n{args}"
        if args.size  2 then throwError "Recursive call with invalid number of parameters: {args}"
        let i_arg := args.get! 0
        let x_arg := args.get! 1
        let eIsValid  mkAppM ``FixI.is_valid_p_rec #[k_var, i_arg, x_arg]
        trace[Diverge.def.valid] "rec: result: \n{eIsValid}"
        pure eIsValid
      else do
        -- Remaining case: normal application.
        -- It shouldn't use the continuation.
        -- TODO: actually, it can
        proveNoKExprIsValid k_var e

-- Prove that a match expression is valid.
partial def proveMatchIsValid (k_var kk_var : Expr) (me : MatchInfo) : MetaM Expr := do
  trace[Diverge.def.valid] "proveMatchIsValid: {me}"
  -- Prove the validity of the branch expressions
  let branchesValid:Array Expr  me.branches.mapIdxM fun idx br => do
    -- Go inside the lambdas - note that we have to be careful: some of the
    -- binders might come from the match, and some of the binders might come
    -- from the fact that the expression in the match is a lambda expression:
    -- we use the branchesNumParams field for this reason
    let numParams := me.branchesNumParams.get! idx
    lambdaTelescopeN br numParams fun xs br => do
    -- Prove that the branch expression is valid
    let brValid  proveExprIsValid k_var kk_var br
    -- Reconstruct the lambda expression
    mkLambdaFVars xs brValid
  trace[Diverge.def.valid] "branchesValid:\n{branchesValid}"
  -- Compute the motive, which has the following shape:
  -- ```
  -- λ scrut => is_valid_p k (λ k => match scrut with ...)
  --                                 ^^^^^^^^^^^^^^^^^^^^
  --                         this is the original match expression, with the
  --                        the difference that the scrutinee(s) is a variable
  -- ```
  let validMotive : Expr  do
    -- The motive is a function of the scrutinees (i.e., a lambda expression):
    -- introduce binders for the scrutinees
    let declInfos := me.scruts.mapIdx fun idx scrut =>
      let name : Name := mkAnonymous "scrut" idx
      let ty := λ (_ : Array Expr) => inferType scrut
      (name, ty)
    withLocalDeclsD declInfos fun scrutVars => do
    -- Create a match expression but where the scrutinees have been replaced
    -- by variables
    let params : Array (Option Expr) := me.params.map some
    let motive : Option Expr := some me.motive
    let scruts : Array (Option Expr) := scrutVars.map some
    let branches : Array (Option Expr) := me.branches.map some
    let args := params ++ [motive] ++ scruts ++ branches
    let matchE  mkAppOptM me.matcherName args
    -- Wrap in the `is_valid_p` predicate
    let matchE  mkLambdaFVars #[kk_var] matchE
    let validMotive  mkAppM ``FixI.is_valid_p #[k_var, matchE]
    -- Abstract away the scrutinee variables
    mkLambdaFVars scrutVars validMotive
  trace[Diverge.def.valid] "valid motive: {validMotive}"
  -- Put together
  let valid  do
    -- We let Lean infer the parameters
    let params : Array (Option Expr) := me.params.map (λ _ => none)
    let motive := some validMotive
    let scruts := me.scruts.map some
    let branches := branchesValid.map some
    let args := params ++ [motive] ++ scruts ++ branches
    mkAppOptM me.matcherName args
  trace[Diverge.def.valid] "proveMatchIsValid:\n{valid}:\n{← inferType valid}"
  pure valid

end

-- Prove that a single body (in the mutually recursive group) is valid.
--
-- For instance, if we define the mutually recursive group [`is_even`, `is_odd`],
-- we prove that `is_even.body` and `is_odd.body` are valid.
partial def proveSingleBodyIsValid
  (k_var : Expr) (preDef : PreDefinition) (bodyConst : Expr) :
  MetaM Expr := do
  trace[Diverge.def.valid] "proveSingleBodyIsValid: bodyConst: {bodyConst}"
  -- Lookup the definition (`bodyConst` is a const, we want to retrieve its
  -- definition to dive inside)
  let name := bodyConst.constName!
  let env  getEnv
  let body := (env.constants.find! name).value!
  trace[Diverge.def.valid] "body: {body}"
  lambdaTelescope body fun xs body => do
  assert! xs.size = 2
  let kk_var := xs.get! 0
  let x_var := xs.get! 1
  -- State the type of the theorem to prove
  let thmTy  mkAppM ``FixI.is_valid_p
    #[k_var,  mkLambdaFVars #[kk_var] ( mkAppM' bodyConst #[kk_var, x_var])]
  trace[Diverge.def.valid] "thmTy: {thmTy}"
  -- Prove that the body is valid
  let proof  proveExprIsValid k_var kk_var body
  let proof  mkLambdaFVars #[k_var, x_var] proof
  trace[Diverge.def.valid] "proveSingleBodyIsValid: proof:\n{proof}:\n{← inferType proof}"
  -- The target type (we don't have to do this: this is simply a sanity check,
  -- and this allows a nicer debugging output)
  let thmTy  do
    let body  mkAppM' bodyConst #[kk_var, x_var]
    let body  mkLambdaFVars #[kk_var] body
    let ty  mkAppM ``FixI.is_valid_p #[k_var, body]
    mkForallFVars #[k_var, x_var] ty
  trace[Diverge.def.valid] "proveSingleBodyIsValid: thmTy\n{thmTy}:\n{← inferType thmTy}"
  -- Save the theorem
  let name := preDef.declName ++ "body_is_valid"
  let decl := Declaration.thmDecl {
    name
    levelParams := preDef.levelParams
    type := thmTy
    value := proof
    all := [name]
  }
  addDecl decl
  trace[Diverge.def.valid] "proveSingleBodyIsValid: added thm: {name}"
  -- Return the theorem
  pure (Expr.const name (preDef.levelParams.map .param))

-- Prove that the list of bodies are valid.
--
-- For instance, if we define the mutually recursive group [`is_even`, `is_odd`],
-- we prove that `Funs.Cons is_even.body (Funs.Cons is_odd.body Funs.Nil)` is
-- valid.
partial def proveFunsBodyIsValid (inOutTys: Expr) (bodyFuns : Expr)
  (k_var : Expr) (bodiesValid : Array Expr) : MetaM Expr := do
  -- Create the big "and" expression, which groups the validity proof of the individual bodies
  let rec mkValidConj (i : Nat) : MetaM Expr := do
    if i = bodiesValid.size then
      -- We reached the end
      mkAppM ``FixI.Funs.is_valid_p_Nil #[k_var]
    else do
      -- We haven't reached the end: introduce a conjunction
      let valid := bodiesValid.get! i
      let valid  mkAppM' valid #[k_var]
      mkAppM ``And.intro #[valid,  mkValidConj (i + 1)]
  let andExpr  mkValidConj 0
  -- Wrap in the `is_valid_p_is_valid_p` theorem, and abstract the continuation
  let isValid  mkAppM ``FixI.Funs.is_valid_p_is_valid_p #[inOutTys, k_var, bodyFuns, andExpr]
  mkLambdaFVars #[k_var] isValid

-- Prove that the mut rec body (i.e., the unary body which groups the bodies
-- of all the functions in the mutually recursive group and on which we will
-- apply the fixed-point operator) is valid.
--
-- We save the proof in the theorem "[GROUP_NAME]."mut_rec_body_is_valid",
-- which we return.
--
-- TODO: maybe this function should introduce k_var itself
def proveMutRecIsValid
  (grName : Name) (grLvlParams : List Name)
  (inOutTys : Expr) (bodyFuns mutRecBodyConst : Expr)
  (k_var : Expr) (preDefs : Array PreDefinition)
  (bodies : Array Expr) : MetaM Expr := do
  -- First prove that the individual bodies are valid
  let bodiesValid 
    bodies.mapIdxM fun idx body => do
      let preDef := preDefs.get! idx
      trace[Diverge.def.valid] "## Proving that the body {body} is valid"
      proveSingleBodyIsValid k_var preDef body
  -- Then prove that the mut rec body is valid
  trace[Diverge.def.valid] "## Proving that the 'Funs' body is valid"
  let isValid  proveFunsBodyIsValid inOutTys bodyFuns k_var bodiesValid
  -- Save the theorem
  let thmTy  mkAppM ``FixI.is_valid #[mutRecBodyConst]
  let name := grName ++ "mut_rec_body_is_valid"
  let decl := Declaration.thmDecl {
    name
    levelParams := grLvlParams
    type := thmTy
    value := isValid
    all := [name]
  }
  addDecl decl
  trace[Diverge.def.valid] "proveFunsBodyIsValid: added thm: {name}:\n{thmTy}"
  -- Return the theorem
  pure (Expr.const name (grLvlParams.map .param))

-- Generate the final definions by using the mutual body and the fixed point operator.
--
-- For instance:
-- ```
-- def is_even (i : Int) : Result Bool := mut_rec_body 0 i
-- def is_odd (i : Int) : Result Bool := mut_rec_body 1 i
-- ```
def mkDeclareFixDefs (mutRecBody : Expr) (inOutTys : Array (Expr × Expr)) (preDefs : Array PreDefinition) :
  TermElabM (Array Name) := do
  let grSize := preDefs.size
  let defs  preDefs.mapIdxM fun idx preDef => do
    lambdaTelescope preDef.value fun xs _ => do
    -- Retrieve the input type
    let in_ty := (inOutTys.get! idx.val).fst
    -- Create the index
    let idx  mkFinVal grSize idx.val
    -- Group the inputs into a dependent tuple
    let input  mkSigmasVal in_ty xs.toList
    -- Apply the fixed point
    let fixedBody  mkAppM ``FixI.fix #[mutRecBody, idx, input]
    let fixedBody  mkLambdaFVars xs fixedBody
    -- Create the declaration
    let name := preDef.declName
    let decl := Declaration.defnDecl {
      name := name
      levelParams := preDef.levelParams
      type := preDef.type
      value := fixedBody
      hints := ReducibilityHints.regular (getMaxHeight ( getEnv) fixedBody + 1)
      safety := .safe
      all := [name]
    }
    addDecl decl
    pure name
  pure defs

-- Prove the equations that we will use as unfolding theorems
partial def proveUnfoldingThms (isValidThm : Expr) (inOutTys : Array (Expr × Expr))
  (preDefs : Array PreDefinition) (decls : Array Name) : MetaM Unit := do
  let grSize := preDefs.size
  let proveIdx (i : Nat) : MetaM Unit := do
    let preDef := preDefs.get! i
    let defName := decls.get! i
    -- Retrieve the arguments
    lambdaTelescope preDef.value fun xs body => do
    trace[Diverge.def.unfold] "proveUnfoldingThms: xs: {xs}"
    trace[Diverge.def.unfold] "proveUnfoldingThms: body: {body}"
    -- The theorem statement
    let thmTy  do
      -- The equation: the declaration gives the lhs, the pre-def gives the rhs
      let lhs  mkAppOptM defName (xs.map some)
      let rhs := body
      let eq  mkAppM ``Eq #[lhs, rhs]
      mkForallFVars xs eq
    trace[Diverge.def.unfold] "proveUnfoldingThms: thm statement: {thmTy}"
    -- The proof
    -- Use the fixed-point equation
    let proof  mkAppM ``FixI.is_valid_fix_fixed_eq #[isValidThm]
    -- Add the index
    let idx  mkFinVal grSize i
    let proof  mkAppM ``congr_fun #[proof, idx]
    -- Add the input argument
    let arg  mkSigmasVal (inOutTys.get! i).fst xs.toList
    let proof  mkAppM ``congr_fun #[proof, arg]
    -- Abstract the arguments away
    let proof  mkLambdaFVars xs proof
    trace[Diverge.def.unfold] "proveUnfoldingThms: proof: {proof}:\n{← inferType proof}"
    -- Declare the theorem
    let name := preDef.declName ++ "unfold"
    let decl := Declaration.thmDecl {
      name
      levelParams := preDef.levelParams
      type := thmTy
      value := proof
      all := [name]
    }
    addDecl decl
    -- Add the unfolding theorem to the equation compiler
    eqnsAttribute.add preDef.declName #[name]
    trace[Diverge.def.unfold] "proveUnfoldingThms: added thm: {name}:\n{thmTy}"
  let rec prove (i : Nat) : MetaM Unit := do
    if i = preDefs.size then pure ()
    else do
      proveIdx i
      prove (i + 1)
  --
  prove 0

def divRecursion (preDefs : Array PreDefinition) : TermElabM Unit := do
  let msg := toMessageData <| preDefs.map fun pd => (pd.declName, pd.levelParams, pd.type, pd.value)
  trace[Diverge.def] ("divRecursion: defs:\n" ++ msg)

  -- TODO: what is this?
  for preDef in preDefs do
    applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation

  -- Retrieve the name of the first definition, that we will use as the namespace
  -- for the definitions common to the group
  let def0 := preDefs[0]!
  let grName := def0.declName
  trace[Diverge.def] "group name: {grName}"

  /- # Compute the input/output types of the continuation `k`. -/
  let grLvlParams := def0.levelParams
  trace[Diverge.def] "def0 universe levels: {def0.levelParams}"

  -- We first compute the list of pairs: (input type × output type)
  let inOutTys : Array (Expr × Expr) 
      preDefs.mapM (fun preDef => do
        withRef preDef.ref do -- is the withRef useful?
        -- Check the universe parameters - TODO: I'm not sure what the best thing
        -- to do is. In practice, all the type parameters should be in Type 0, so
        -- we shouldn't have universe issues.
        if preDef.levelParams  grLvlParams then
          throwError "Non-uniform polymorphism in the universes"
        forallTelescope preDef.type (fun in_tys out_ty => do
          let in_ty  liftM (mkSigmasType in_tys.toList)
          -- Retrieve the type in the "Result"
          let out_ty  getResultTy out_ty
          let out_ty  liftM (mkSigmasMatch in_tys.toList out_ty)
          pure (in_ty, out_ty)
        )
      )
  trace[Diverge.def] "inOutTys: {inOutTys}"
  -- Turn the list of input/output type pairs into an expresion
  let inOutTysExpr  inOutTys.mapM (λ (x, y) => mkInOutTy x y)
  let inOutTysExpr  mkListLit ( inferType (inOutTysExpr.get! 0)) inOutTysExpr.toList

  -- From the list of pairs of input/output types, actually compute the
  -- type of the continuation `k`.
  -- We first introduce the index `i : Fin n` where `n` is the number of
  -- functions in the group.
  let i_var_ty := mkFin preDefs.size
  withLocalDeclD (mkAnonymous "i" 0) i_var_ty fun i_var => do
  let in_out_ty  mkAppM ``List.get #[inOutTysExpr, i_var]
  trace[Diverge.def] "in_out_ty := {in_out_ty} : {← inferType in_out_ty}"
  -- Add an auxiliary definition for `in_out_ty`
  let in_out_ty  do
    let value  mkLambdaFVars #[i_var] in_out_ty
    let name := grName.append "in_out_ty"
    let levelParams := grLvlParams
    let decl := Declaration.defnDecl {
      name := name
      levelParams := levelParams
      type :=  inferType value
      value := value
      hints := .abbrev
      safety := .safe
      all := [name]
    }
    addDecl decl
    -- Return the constant
    let in_out_ty := Lean.mkConst name (levelParams.map .param)
    mkAppM' in_out_ty #[i_var]
  trace[Diverge.def] "in_out_ty (after decl) := {in_out_ty} : {← inferType in_out_ty}"
  let in_ty  mkAppM ``Sigma.fst #[in_out_ty]
  trace[Diverge.def] "in_ty: {in_ty}"
  withLocalDeclD (mkAnonymous "x" 1) in_ty fun input => do
  let out_ty  mkAppM' ( mkAppM ``Sigma.snd #[in_out_ty]) #[input]
  trace[Diverge.def] "out_ty: {out_ty}"

  -- Introduce the continuation `k`
  let in_ty  mkLambdaFVars #[i_var] in_ty
  let out_ty  mkLambdaFVars #[i_var, input] out_ty
  let kk_var_ty  mkAppM ``FixI.kk_ty #[i_var_ty, in_ty, out_ty]
  trace[Diverge.def] "kk_var_ty: {kk_var_ty}"
  withLocalDeclD (mkAnonymous "kk" 2) kk_var_ty fun kk_var => do
  trace[Diverge.def] "kk_var: {kk_var}"

  -- Replace the recursive calls in all the function bodies by calls to the
  -- continuation `k` and and generate for those bodies declarations
  trace[Diverge.def] "# Generating the unary bodies"
  let bodies  mkDeclareUnaryBodies grLvlParams kk_var inOutTys preDefs
  trace[Diverge.def] "Unary bodies (after decl): {bodies}"
  -- Generate the mutually recursive body
  trace[Diverge.def] "# Generating  the mut rec body"
  let (bodyFuns, mutRecBody)  mkDeclareMutRecBody grName grLvlParams kk_var i_var in_ty out_ty inOutTys.toList bodies
  trace[Diverge.def] "mut rec body (after decl): {mutRecBody}"

  -- Prove that the mut rec body satisfies the validity criteria required by
  -- our fixed-point
  let k_var_ty  mkAppM ``FixI.k_ty #[i_var_ty, in_ty, out_ty]
  withLocalDeclD (mkAnonymous "k" 3) k_var_ty fun k_var => do
  trace[Diverge.def] "# Proving that the mut rec body is valid"
  let isValidThm  proveMutRecIsValid grName grLvlParams inOutTysExpr bodyFuns mutRecBody k_var preDefs bodies

  -- Generate the final definitions
  trace[Diverge.def] "# Generating the final definitions"
  let decls  mkDeclareFixDefs mutRecBody inOutTys preDefs

  -- Prove the unfolding theorems
  trace[Diverge.def] "# Proving the unfolding theorems"
  proveUnfoldingThms isValidThm inOutTys preDefs decls

  -- Generating code -- TODO
  addAndCompilePartialRec preDefs

-- The following function is copy&pasted from Lean.Elab.PreDefinition.Main
-- This is the only part where we make actual changes and hook into the equation compiler.
-- (I've removed all the well-founded stuff to make it easier to read though.)

open private ensureNoUnassignedMVarsAtPreDef betaReduceLetRecApps partitionPreDefs
  addAndCompilePartial addAsAxioms from Lean.Elab.PreDefinition.Main

def addPreDefinitions (preDefs : Array PreDefinition) : TermElabM Unit := withLCtx {} {} do
  for preDef in preDefs do
    trace[Diverge.elab] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
  let preDefs  preDefs.mapM ensureNoUnassignedMVarsAtPreDef
  let preDefs  betaReduceLetRecApps preDefs
  let cliques := partitionPreDefs preDefs
  let mut hasErrors := false
  for preDefs in cliques do
    trace[Diverge.elab] "{preDefs.map (·.declName)}"
    try
      withRef (preDefs[0]!.ref) do
        divRecursion preDefs
    catch ex =>
      -- If it failed, we add the functions as partial functions
      hasErrors := true
      logException ex
      let s  saveState
      try
        if preDefs.all fun preDef => preDef.kind == DefKind.def ||
           preDefs.all fun preDef => preDef.kind == DefKind.abbrev then
          -- try to add as partial definition
          try
            addAndCompilePartial preDefs (useSorry := true)
          catch _ =>
            -- Compilation failed try again just as axiom
            s.restore
            addAsAxioms preDefs
        else return ()
      catch _ => s.restore

-- The following two functions are copy-pasted from Lean.Elab.MutualDef

open private elabHeaders levelMVarToParamHeaders getAllUserLevelNames withFunLocalDecls elabFunValues
  instantiateMVarsAtHeader instantiateMVarsAtLetRecToLift checkLetRecsToLiftTypes withUsed from Lean.Elab.MutualDef

def Term.elabMutualDef (vars : Array Expr) (views : Array DefView) : TermElabM Unit := do
    let scopeLevelNames  getLevelNames
    let headers  elabHeaders views
    let headers  levelMVarToParamHeaders views headers
    let allUserLevelNames := getAllUserLevelNames headers
    withFunLocalDecls headers fun funFVars => do
      for view in views, funFVar in funFVars do
        addLocalVarInfo view.declId funFVar
        -- Add fake use site to prevent "unused variable" warning (if the
        -- function is actually not recursive, Lean would print this warning).
        -- Remark: we could detect this case and encode the function without
        -- using the fixed-point. In practice it shouldn't happen however:
        -- we define non-recursive functions with the `divergent` keyword
        -- only for testing purposes.
        addTermInfo' view.declId funFVar
      let values 
        try
          let values  elabFunValues headers
          Term.synthesizeSyntheticMVarsNoPostponing
          values.mapM (instantiateMVars ·)
        catch ex =>
          logException ex
          headers.mapM fun header => mkSorry header.type (synthetic := true)
      let headers  headers.mapM instantiateMVarsAtHeader
      let letRecsToLift  getLetRecsToLift
      let letRecsToLift  letRecsToLift.mapM instantiateMVarsAtLetRecToLift
      checkLetRecsToLiftTypes funFVars letRecsToLift
      withUsed vars headers values letRecsToLift fun vars => do
        let preDefs  MutualClosure.main vars headers funFVars values letRecsToLift
        for preDef in preDefs do
          trace[Diverge.elab] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
        let preDefs  withLevelNames allUserLevelNames <| levelMVarToParamPreDecls preDefs
        let preDefs  instantiateMVarsAtPreDecls preDefs
        let preDefs  fixLevelParams preDefs scopeLevelNames allUserLevelNames
        for preDef in preDefs do
          trace[Diverge.elab] "after eraseAuxDiscr, {preDef.declName} : {preDef.type} :=\n{preDef.value}"
        checkForHiddenUnivLevels allUserLevelNames preDefs
        addPreDefinitions preDefs

open Command in
def Command.elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
  let views  ds.mapM fun d => do
    let `($mods:declModifiers divergent def $id:declId $sig:optDeclSig $val:declVal) := d
      | throwUnsupportedSyntax
    let modifiers  elabModifiers mods
    let (binders, type) := expandOptDeclSig sig
    let deriving? := none
    pure { ref := d, kind := DefKind.def, modifiers,
           declId := id, binders, type? := type, value := val, deriving? }
  runTermElabM fun vars => Term.elabMutualDef vars views

-- Special command so that we don't fall back to the built-in mutual when we produce an error.
local syntax "_divergent" Parser.Command.mutual : command
elab_rules : command | `(_divergent mutual $decls* end) => Command.elabMutualDef decls

macro_rules
  | `(mutual $decls* end) => do
    unless !decls.isEmpty && decls.all (·.1.getKind == ``divergentDef) do
      Macro.throwUnsupported
    `(command| _divergent mutual $decls* end)

open private setDeclIdName from Lean.Elab.Declaration
elab_rules : command
  | `($mods:declModifiers divergent%$tk def $id:declId $sig:optDeclSig $val:declVal) => do
    let (name, _) := expandDeclIdCore id
    if (`_root_).isPrefixOf name then throwUnsupportedSyntax
    let view := extractMacroScopes name
    let .str ns shortName := view.name | throwUnsupportedSyntax
    let shortName' := { view with name := shortName }.review
    let cmd  `(mutual $mods:declModifiers divergent%$tk def $(setDeclIdName id shortName'):declId $sig:optDeclSig $val:declVal end)
    if ns matches .anonymous then
      Command.elabCommand cmd
    else
      Command.elabCommand <|  `(namespace $(mkIdentFrom id ns) $cmd end $(mkIdentFrom id ns))

namespace Tests
  /- Some examples of partial functions -/

  divergent def list_nth {a: Type} (ls : List a) (i : Int) : Result a :=
    match ls with
    | [] => .fail .panic
    | x :: ls =>
      if i = 0 then return x
      else return ( list_nth ls (i - 1))

  #check list_nth.unfold

  example {a: Type} (ls : List a) :
     (i : Int),
    0  i  i < ls.length 
     x, list_nth ls i = .ret x := by
    induction ls
    . intro i hpos h; simp at h; linarith
    . rename_i hd tl ih
      intro i hpos h
      -- We can directly use `rw [list_nth]`
      rw [list_nth]; simp
      split <;> try simp [*]
      . tauto
      . -- We don't have to do this if we use scalar_tac
        have hneq : 0 < i := by cases i <;> rename_i a _ <;> simp_all; cases a <;> simp_all
        simp at h
        have  x, ih  := ih (i - 1) (by linarith) (by linarith)
        simp [ih]
        tauto

  -- Return a continuation
  divergent def list_nth_with_back {a: Type} (ls : List a) (i : Int) :
    Result (a × (a  Result (List a))) :=
    match ls with
    | [] => .fail .panic
    | x :: ls =>
      if i = 0 then return (x, (λ ret => return (ret :: ls)))
      else do
        let (x, back)  list_nth_with_back ls (i - 1)
        return (x,
          (λ ret => do
           let ls  back ret
           return (x :: ls)))

  #check list_nth_with_back.unfold

  mutual
    divergent def is_even (i : Int) : Result Bool :=
      if i = 0 then return true else return ( is_odd (i - 1))

    divergent def is_odd (i : Int) : Result Bool :=
      if i = 0 then return false else return ( is_even (i - 1))
  end

  #check is_even.unfold
  #check is_odd.unfold

  mutual
    divergent def foo (i : Int) : Result Nat :=
      if i > 10 then return ( foo (i / 10)) + ( bar i) else bar 10

    divergent def bar (i : Int) : Result Nat :=
      if i > 20 then foo (i / 20) else .ret 42
  end

  #check foo.unfold
  #check bar.unfold

  -- Testing dependent branching and let-bindings
  divergent def isNonZero (i : Int) : Result Bool :=
    if _h:i = 0 then return false
    else
      let b := true
      return b

  #check isNonZero.unfold

  -- Testing let-bindings
  divergent def iInBounds {a : Type} (ls : List a) (i : Int) : Result Bool :=
    let i0 := ls.length
    if i < i0
    then Result.ret True
    else Result.ret False

  #check iInBounds.unfold

  divergent def isCons
    {a : Type} (ls : List a) : Result Bool :=
    let ls1 := ls
    match ls1 with
    | [] => Result.ret False
    | _ :: _ => Result.ret True

  #check isCons.unfold

  -- Testing what happens when we use concrete arguments in dependent tuples
  divergent def test1
    (_ : Option Bool) (_ : Unit) :
    Result Unit
    :=
    test1 Option.none ()

  #check test1.unfold

end Tests

end Diverge