lean4-htt/src/Lean/Class.lean

/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Attributes

namespace Lean

inductive ClassEntry :=
  | «class»    (name : Name) (hasOutParam : Bool)
  | «instance» (name : Name) (ofClass : Name) -- TODO: remove after we remove old type class resolution

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) := SMap.empty) -- TODO: delete
  (hasOutParam      : SMap Name Bool := SMap.empty)        -- We should keep only this one
  (instances        : SMap Name Unit := SMap.empty)        -- TODO: delete

namespace ClassState

instance : Inhabited ClassState := ⟨{}⟩

def addEntry (s : ClassState) (entry : ClassEntry) : ClassState :=
  match entry with
  | ClassEntry.«class» clsName hasOutParam =>
    { s with hasOutParam := s.hasOutParam.insert clsName hasOutParam }
  | ClassEntry.«instance» instName clsName =>
    { s with
      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] }

def switch : ClassState → ClassState
  | ⟨m₁, m₂, m₃⟩ => ⟨m₁.switch, m₂.switch, m₃.switch⟩

end ClassState

/- TODO: add support for scoped instances -/
builtin_initialize classExtension : SimplePersistentEnvExtension ClassEntry ClassState ←
  registerSimplePersistentEnvExtension {
    name          := `classExt,
    addEntryFn    := ClassState.addEntry,
    addImportedFn := fun es => (mkStateFromImportedEntries ClassState.addEntry {} es).switch
  }

@[export lean_is_class]
def isClass (env : Environment) (n : Name) : Bool :=
  (classExtension.getState env).hasOutParam.contains n

@[export lean_is_instance]
def isInstance (env : Environment) (n : Name) : Bool :=
  (classExtension.getState env).instances.contains n

@[export lean_get_class_instances]
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]
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]
def isOutParam (e : Expr) : Bool :=
  e.isAppOfArity `outParam 1

/--
  Auxiliary function for checking whether a class has `outParam`, and
  whether they are being correctly used.
  A regular (i.e., non `outParam`) must not depend on an `outParam`.
  Reason for this restriction:
  When performing type class resolution, we replace arguments that
  are `outParam`s with fresh metavariables. If regular parameters could
  depend on `outParam`s, then we would also have to replace them with
  fresh metavariables. Otherwise, the resulting expression could be type
  incorrect. This transformation would be counterintuitive to users since
  we would implicitly treat these regular parameters as `outParam`s.
-/
private partial def checkOutParam : Nat → Array FVarId → Expr → Except String Bool
  | i, outParams, Expr.forallE _ d b _ =>
    if isOutParam d then
      let fvarId    := Name.mkNum `_fvar outParams.size
      let outParams := outParams.push fvarId
      let fvar      := mkFVar fvarId
      let b         := b.instantiate1 fvar
      checkOutParam (i+1) outParams b
    else if d.hasAnyFVar fun fvarId => outParams.contains fvarId then
      Except.error s!"invalid class, parameter #{i} depends on `outParam`, but it is not an `outParam`"
    else
      checkOutParam (i+1) outParams b
  | i, outParams, e => pure (outParams.size > 0)

def addClass (env : Environment) (clsName : Name) : Except String Environment :=
  if isClass env clsName then
    Except.error s!"class has already been declared '{clsName}'"
  else match env.find? clsName with
    | none      => Except.error ("unknown declaration '" ++ toString clsName ++ "'")
    | some decl@(ConstantInfo.inductInfo _) => do
      let b ← checkOutParam 1 #[] decl.type
      Except.ok (classExtension.addEntry env (ClassEntry.«class» clsName b))
    | some _    => Except.error ("invalid 'class', declaration '" ++ toString clsName ++ "' must be inductive datatype or structure")

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.forallE _ _ b _ => getClassName env b
  | e                    => do
    let Expr.const c _ _ ← pure e.getAppFn | none
    let info ← env.find? c
    match info.value? with
    | some val => do
      let body ← consumeNLambdas e.getAppNumArgs val
      getClassName env body
    | none =>
      if isClass env c then some c
      else none

builtin_initialize
  registerBuiltinAttribute {
    name  := `class,
    descr := "type class",
    add   := fun decl args persistent => do
      let env ← getEnv
      if args.hasArgs then throwError "invalid attribute 'class', unexpected argument"
      unless persistent do throwError "invalid attribute 'class', must be persistent"
      let env ← ofExcept (addClass env decl)
      setEnv env
  }

-- TODO: delete
@[export lean_add_instance_old]
def addGlobalInstanceOld (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))

end Lean