From ef154ec0cf8d56e5723169d64e5c7b460ca85eec Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 14 Mar 2022 11:52:54 -0700
Subject: [PATCH] test: add `TreeNode` example using `for in` notation

---
 tests/lean/run/treeNode.lean | 41 ++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 tests/lean/run/treeNode.lean

diff --git a/tests/lean/run/treeNode.lean b/tests/lean/run/treeNode.lean
new file mode 100644
index 0000000000..77c9bb0090
--- /dev/null
+++ b/tests/lean/run/treeNode.lean
@@ -0,0 +1,41 @@
+inductive TreeNode :=
+ | mkLeaf (name : String) : TreeNode
+ | mkNode (name : String) (children : List TreeNode) : TreeNode
+
+def treeToList (t : TreeNode) : List String :=
+ match t with
+ | .mkLeaf name => [name]
+ | .mkNode name children => Id.run do
+   let mut r := [name]
+   for h : child in children do
+     -- We will not this the following `have` in the future
+     have : sizeOf child < 1 + sizeOf name + sizeOf children := Nat.lt_trans (List.sizeOf_lt_of_mem h) (by simp_arith)
+     r := r ++ treeToList child
+   return r
+
+@[simp] theorem treeToList_eq (name : String) (children : List TreeNode) : treeToList (.mkNode name children) =  name :: List.join (children.map treeToList) := by
+  simp [treeToList, Id.run, forIn, List.forIn]
+  have : ∀ acc, (Id.run do List.forIn.loop (fun a b => ForInStep.yield (b ++ treeToList a)) children acc) = acc ++ List.join (List.map treeToList children) := by
+    intro acc
+    induction children generalizing acc with simp [List.forIn.loop, List.map, List.join, Id.run]
+    | cons c cs ih => simp [Id.run] at ih; simp [ih, List.append_assoc]
+  apply this
+
+mutual
+  def numNames : TreeNode → Nat
+    | .mkLeaf _  => 1
+    | .mkNode _ cs => 1 + numNamesLst cs
+  def numNamesLst : List TreeNode → Nat
+    | [] => 0
+    | a :: as => numNames a + numNamesLst as
+end
+
+theorem length_treeToList_eq_numNames (t : TreeNode) : (treeToList t).length = numNames t := by
+  match t with
+  | .mkLeaf .. => simp [treeToList, numNames]
+  | .mkNode _ cs => simp_arith [numNames, helper cs]
+where
+  helper (cs : List TreeNode) : (cs.map treeToList).join.length = numNamesLst cs := by
+    match cs with
+    | [] => rfl
+    | c::cs' => simp [List.join, List.map, numNamesLst, length_treeToList_eq_numNames c, helper cs']