diff --git a/DocGen4/Hierarchy.lean b/DocGen4/Hierarchy.lean new file mode 100644 index 0000000..d023e85 --- /dev/null +++ b/DocGen4/Hierarchy.lean @@ -0,0 +1,72 @@ +import Lean +import Std.Data.HashMap + +namespace DocGen4 + +open Lean Std Name + +def getDepth : Name → Nat +| Name.anonymous => 0 +| Name.str p _ _ => getDepth p + 1 +| Name.num p _ _ => getDepth p + 1 + +def getNLevels (name : Name) (levels: Nat) : Name := + (components.drop (components.length - levels)).reverse.foldl (· ++ ·) Name.anonymous + where + components := name.components' + +inductive Hierarchy where +| node : Name → RBNode Name (λ _ => Hierarchy) → Hierarchy + +instance : Inhabited Hierarchy := ⟨Hierarchy.node Name.anonymous RBNode.leaf⟩ + +abbrev HierarchyMap := RBNode Name (λ _ => Hierarchy) + +-- Everything in this namespace is adapted from stdlib's RBNode +namespace HierarchyMap + +def toList : HierarchyMap → List (Name × Hierarchy) +| t => t.revFold (fun ps k v => (k, v)::ps) [] + +def hForIn [Monad m] (t : HierarchyMap) (init : σ) (f : (Name × Hierarchy) → σ → m (ForInStep σ)) : m σ := + t.forIn init (fun a b acc => f (a, b) acc) + +instance : ForIn m HierarchyMap (Name × Hierarchy) where + forIn := HierarchyMap.hForIn + +end HierarchyMap + +namespace Hierarchy + +def empty (n : Name) : Hierarchy := node n RBNode.leaf + +def getName : Hierarchy → Name +| node n _ => n + +def getChildren : Hierarchy → HierarchyMap +| node _ c => c + +partial def insert! (h : Hierarchy) (n : Name) : Hierarchy := do + let hn := h.getName + let mut cs := h.getChildren + if getDepth hn ≥ getDepth n then + panic! "Invalid insert" + else if getDepth hn + 1 == getDepth n then + match cs.find Name.cmp n with + | none => + node hn (cs.insert Name.cmp n $ empty n) + | some _ => h + else + let leveledName := getNLevels n (getDepth hn + 1) + match cs.find Name.cmp leveledName with + | some nextLevel => + cs := cs.erase Name.cmp leveledName + node hn $ cs.insert Name.cmp leveledName (nextLevel.insert! n) + | none => + let child := (insert! (empty leveledName) n) + node hn $ cs.insert Name.cmp leveledName child + +partial def fromArray (names : Array Name) : Hierarchy := names.foldl insert! (empty anonymous) + +end Hierarchy +end DocGen4