From 1a40a38bf1d67ae4c2b319958d616cf7e2cafd13 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 26 Jun 2019 19:43:12 -0700
Subject: [PATCH] feat(library/init/lean/class): add helper functions

---
 library/init/lean/class.lean | 47 ++++++++++++++++++++++++++++++++----
 library/init/lean/expr.lean  | 18 +++++++++++---
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/library/init/lean/class.lean b/library/init/lean/class.lean
index 070447e997..f9148ef636 100644
--- a/library/init/lean/class.lean
+++ b/library/init/lean/class.lean
@@ -26,24 +26,61 @@ end ClassEntry
 structure ClassState :=
 (classToInstances : SMap Name (List Name) Name.quickLt := SMap.empty)
 (hasOutParam      : SMap Name Bool Name.quickLt := SMap.empty)
-(instances        : NameSet := ∅)
+(instances        : SMap Name Unit Name.quickLt := SMap.empty)
 
-def ClassState.addEntry (imported : Bool) (s : ClassState) (entry : ClassEntry) : ClassState :=
+namespace ClassState
+
+instance : Inhabited ClassState := ⟨{}⟩
+
+def addEntry (s : ClassState) (entry : ClassEntry) : ClassState :=
 match entry with
 | ClassEntry.«class» clsName hasOutParam :=
   { hasOutParam := s.hasOutParam.insert clsName hasOutParam, .. s }
 | ClassEntry.«instance» instName clsName :=
-  { instances        := if imported then s.instances else s.instances.insert instName,
+  { instances        := s.instances.insert instName (),
     classToInstances := match s.classToInstances.find clsName with
       | some insts := s.classToInstances.insert clsName (instName :: insts)
       | none       := s.classToInstances.insert clsName [instName],
     .. s }
 
+end ClassState
+
+/- TODO: add support for scoped instances -/
 def mkClassExtension : IO (SimplePersistentEnvExtension ClassEntry ClassState) :=
 registerSimplePersistentEnvExtension {
   name          := `class,
-  addEntryFn    := ClassState.addEntry false,
-  addImportedFn := mkStateFromImportedEntries (ClassState.addEntry true) {}
+  addEntryFn    := ClassState.addEntry,
+  addImportedFn := mkStateFromImportedEntries ClassState.addEntry {}
 }
 
+@[init mkClassExtension]
+constant classExtension : SimplePersistentEnvExtension ClassEntry ClassState := default _
+
+def isClass (env : Environment) (n : Name) : Bool :=
+(classExtension.getState env).hasOutParam.contains n
+
+def isInstance (env : Environment) (n : Name) : Bool :=
+(classExtension.getState env).instances.contains n
+
+def getClassInstances (env : Environment) (n : Name) : List Name :=
+match (classExtension.getState env).classToInstances.find n with
+| some insts := insts
+| none       := []
+
+private def isOutParam (e : Expr) : Bool :=
+e.isAppOfArity `outParam 1
+
+private def hasOutParam : Expr → Bool
+| (Expr.pi _ _ d b) := isOutParam d || hasOutParam b
+| _                 := false
+
+def addClass (env : Environment) (clsName : Name) : Except String Environment :=
+if isClass env clsName then Except.error ("class has already been declared '" ++ toString clsName ++ "'")
+else match env.find clsName with
+  | none      := Except.error ("unknown declaration '" ++ toString clsName ++ "'")
+  | some decl := Except.ok (classExtension.addEntry env (ClassEntry.«class» clsName (hasOutParam decl.type)))
+
+-- def addInstance (env : Environment) (instName : Name) : Environment :=
+-- match env.find
+
 end Lean
diff --git a/library/init/lean/expr.lean b/library/init/lean/expr.lean
index fb6a539199..fd6f88eaf4 100644
--- a/library/init/lean/expr.lean
+++ b/library/init/lean/expr.lean
@@ -94,15 +94,25 @@ constant eqv (a : @& Expr) (b : @& Expr) : Bool := default _
 
 instance : HasBeq Expr := ⟨Expr.eqv⟩
 
+def getAppFn : Expr → Expr
+| (Expr.app f a) := getAppFn f
+| e              := e
+
+def isAppOf (e : Expr) (n : Name) : Bool :=
+match e.getAppFn with
+| Expr.const c _ := c == n
+| _ := false
+
+def isAppOfArity : Expr → Name → Nat → Bool
+| (Expr.const c _) n 0     := c == n
+| (Expr.app f _)   n (a+1) := isAppOfArity f n a
+| _                _ _     := false
+
 end Expr
 
 def mkConst (n : Name) (ls : List Level := []) : Expr :=
 Expr.const n ls
 
-def getAppFn : Expr → Expr
-| (Expr.app f a) := getAppFn f
-| e              := e
-
 def mkBinApp (f a b : Expr) :=
 Expr.app (Expr.app f a) b