50 lines
No EOL
1.7 KiB
Text
50 lines
No EOL
1.7 KiB
Text
/-
|
|
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
Authors: Leonardo de Moura
|
|
-/
|
|
import Lean.Meta.Basic
|
|
|
|
namespace Lean.Meta
|
|
|
|
namespace SortLocalDecls
|
|
|
|
structure Context where
|
|
localDecls : NameMap LocalDecl := {}
|
|
|
|
structure State where
|
|
visited : NameSet := {}
|
|
result : Array LocalDecl := #[]
|
|
|
|
abbrev M := ReaderT Context $ StateRefT State MetaM
|
|
|
|
mutual
|
|
partial def visitExpr (e : Expr) : M Unit := do
|
|
match e with
|
|
| Expr.proj _ _ e _ => visitExpr e
|
|
| Expr.forallE _ d b _ => visitExpr d; visitExpr b
|
|
| Expr.lam _ d b _ => visitExpr d; visitExpr b
|
|
| Expr.letE _ t v b _ => visitExpr t; visitExpr v; visitExpr b
|
|
| Expr.app f a _ => visitExpr f; visitExpr a
|
|
| Expr.mdata _ b _ => visitExpr b
|
|
| Expr.mvar _ _ => let v ← instantiateMVars e; unless v.isMVar do visitExpr v
|
|
| Expr.fvar fvarId _ => if let some localDecl := (← read).localDecls.find? fvarId then visitLocalDecl localDecl
|
|
| _ => return ()
|
|
|
|
partial def visitLocalDecl (localDecl : LocalDecl) : M Unit := do
|
|
unless (← get).visited.contains localDecl.fvarId do
|
|
modify fun s => { s with visited := s.visited.insert localDecl.fvarId }
|
|
visitExpr localDecl.type
|
|
if let some val := localDecl.value? then
|
|
visitExpr val
|
|
modify fun s => { s with result := s.result.push localDecl }
|
|
end
|
|
|
|
end SortLocalDecls
|
|
|
|
open SortLocalDecls in
|
|
def sortLocalDecls (localDecls : Array LocalDecl) : MetaM (Array LocalDecl) :=
|
|
let aux : M (Array LocalDecl) := do localDecls.forM visitLocalDecl; return (← get).result
|
|
aux.run { localDecls := localDecls.foldl (init := {}) fun s d => s.insert d.fvarId d } |>.run' {}
|
|
|
|
end Lean.Meta |