bookshelf-doc/DocGen4/Hierarchy.lean

78 lines
2.2 KiB
Plaintext
Raw Blame History

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

/-
Copyright (c) 2021 Henrik Böving. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Henrik Böving
-/
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