lean4-htt/library/init/lean/class.lean

143 lines
5 KiB
Text

/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import init.lean.attributes
namespace Lean
inductive ClassEntry
| «class» (name : Name) (hasOutParam : Bool)
| «instance» (name : Name) (ofClass : Name)
namespace ClassEntry
@[inline] def getName : ClassEntry → Name
| «class» n _ => n
| «instance» n _ => n
def lt (a b : ClassEntry) : Bool :=
Name.quickLt a.getName b.getName
end ClassEntry
structure ClassState :=
(classToInstances : SMap Name (List Name) Name.quickLt := SMap.empty)
(hasOutParam : SMap Name Bool Name.quickLt := SMap.empty)
(instances : SMap Name Unit Name.quickLt := SMap.empty)
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 := 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 }
def switch : ClassState → ClassState
| ⟨m₁, m₂, m₃⟩ => ⟨m₁.switch, m₂.switch, m₃.switch⟩
end ClassState
/- TODO: add support for scoped instances -/
def mkClassExtension : IO (SimplePersistentEnvExtension ClassEntry ClassState) :=
registerSimplePersistentEnvExtension {
name := `classExt,
addEntryFn := ClassState.addEntry,
addImportedFn := fun es => (mkStateFromImportedEntries ClassState.addEntry {} es).switch
}
@[init mkClassExtension]
constant classExtension : SimplePersistentEnvExtension ClassEntry ClassState := default _
@[export lean.is_class_core]
def isClass (env : Environment) (n : Name) : Bool :=
(classExtension.getState env).hasOutParam.contains n
@[export lean.is_instance_core]
def isInstance (env : Environment) (n : Name) : Bool :=
(classExtension.getState env).instances.contains n
@[export lean.get_class_instances_core]
def getClassInstances (env : Environment) (n : Name) : List Name :=
match (classExtension.getState env).classToInstances.find n with
| some insts => insts
| none => []
@[export lean.has_out_params_core]
def hasOutParams (env : Environment) (n : Name) : Bool :=
match (classExtension.getState env).hasOutParam.find n with
| some b => b
| none => false
@[export lean.is_out_param_core]
private def isOutParam (e : Expr) : Bool :=
e.isAppOfArity `outParam 1
def Expr.hasOutParam : Expr → Bool
| Expr.pi _ _ d b => isOutParam d || Expr.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 decl.type.hasOutParam))
private def consumeNLambdas : Nat → Expr → Option Expr
| 0, e => some e
| i+1, Expr.lam _ _ _ b => consumeNLambdas i b
| _, _ => none
partial def getClassName (env : Environment) : Expr → Option Name
| Expr.pi _ _ _ d => getClassName d
| e => do
Expr.const c _ ← pure e.getAppFn | none;
info ← env.find c;
match info.value with
| some val => do
body ← consumeNLambdas e.getAppNumArgs val;
getClassName body
| none =>
if isClass env c then some c
else none
@[export lean.add_instance_core]
def addInstance (env : Environment) (instName : Name) : Except String Environment :=
match env.find instName with
| none => Except.error ("unknown declaration '" ++ toString instName ++ "'")
| some decl =>
match getClassName env decl.type with
| none => Except.error ("invalid instance '" ++ toString instName ++ "', failed to retrieve class")
| some clsName => Except.ok (classExtension.addEntry env (ClassEntry.«instance» instName clsName))
@[init] def registerClassAttr : IO Unit :=
registerAttribute {
name := `class,
descr := "type class",
add := fun env decl args persistent => do
unless args.isMissing $ throw (IO.userError ("invalid attribute 'class', unexpected argument"));
unless persistent $ throw (IO.userError ("invalid attribute 'class', must be persistent"));
IO.ofExcept (addClass env decl)
}
@[init] def registerInstanceAttr : IO Unit :=
registerAttribute {
name := `instance,
descr := "type class instance",
add := fun env decl args persistent => do
unless args.isMissing $ throw (IO.userError ("invalid attribute 'instance', unexpected argument"));
unless persistent $ throw (IO.userError ("invalid attribute 'instance', must be persistent"));
IO.ofExcept (addInstance env decl)
}
end Lean