Support for "safety wrapper" structs

Stream: lean4

[Rish Vaishnav]

While programming in Lean, I have found it helpful to make "safety wrapper" types around other types, that give an (informal) "guarantee" that an object satisfies some property. For example, I can make such a wrapper around natural numbers to indicate that the number is even:

inductive N : Type where
| z : N
| s : N → N

/--
A "safety" wrapper around `N`, for function types that expect even number args
and/or output even numbers. The constructor `N2.mk n` should only be used if we
are absolutely sure that `n` is even.
-/
-- @[safety N] -- some annotation that indicates that `N2` is a safety wrapper for `N`
structure N2 : Type where
toN : N

-- it is okay to coerce ONLY in this direction
instance : Coe N2 N := ⟨(N2.toN)⟩

However, say I already have some function f defined on N:

def N.f (n : N) : N := (s (s n))

Even though it is okay to use this function with N2, in order to do so (and preserve typing) I would have to do a double cast along the lines of N2.mk (N.f (N2.toN n)) (with n : N2). To avoid having to do so everywhere, I would probably want to define a function N2.f, but still this is not so convenient because I would also have to do this for every other function defined for N that produces an even number from even numbers.

I think that what I would like is some command that will abstract N.f as follows:

--some command to say that it is okay to use `N.f` with `N2`
safety_ok N.f N2

def N.f' {T : Type} (fTN : T → N) (fNT : N → T) (t : T) : T := fNT (s (s (fTN t)))

and subsequently elaborate usages of N.f n with n : N2 to @N.f' N2 N2.toN N2.mk n. Does anything like this exist?

[Henrik Böving]

No, that said it might be possible to write a meta program that does this somewhat easily. The thing that gets closest to this is probably the coercion mechanism I guess.

[Rish Vaishnav]

Oh, ok. Unfortunately coercion was a bit too "strong" for my purposes, while it's okay to coerce from N2 to N you obviously don't want to coerce in the other direction. I guess what I'm getting at is a kind of partial coercion that only applies in certain cases (that you have to specify).

Anyways, we'll see if I eventually get around to metaprogramming this (if it's as easy as you suggest :sweat_smile:)

[Rish Vaishnav]

To flesh out the idea a bit more, I think we could have a class PCoe for partial coercions for which you can register certain cases where Lean is allowed to insert these coercions:

/--
`PCoe α β` is the typeclass for partial coercions from `α` to `β`. It can be
transitively chained with other `PCoe` instances (all instances of `Coe` are
also instances of `PCoe`). Partial coercion may be used when `x` has type `α`
but it is used in a context where `β` is expected. Whether or not partial
coercion is used depends on a set of user-defined conditions based on the
specific function/constructor constants and argument types being applied.
-/
class PCoe (α : semiOutParam (Sort u)) (β : Sort v) where
  pcoe : α → β

instance [c : Coe α β] : PCoe α β := ⟨c.coe⟩

syntax:1024 (name := pcoeNotation) "⇡" term:1024 : term

@[builtin_term_elab pcoeNotation] def Lean.Elab.Term.elabPCoe
  : Lean.Elab.Term.TermElab := Lean.Elab.Term.elabCoe

Using this, my example above could be rewritten as:

inductive N : Type where
  | z : N
  | s : N → N

/--
A "safety" wrapper around `N`, for function types that expect even number args
and/or output even numbers. The constructor `N2.mk n` should only be used if we
are absolutely sure that `n` is even.
-/
structure N2 : Type where
toN : N

-- we can coerce unconditionally in this direction
instance : Coe N2 N := ⟨N2.toN⟩
-- partial coercion only in this direction
instance : PCoe N N2 := ⟨N2.mk⟩

pcoe_ok N.z : N2
-- `N.z` can be coerced as `⇡N.z : N2`.

def N.f (n : N) : N := s (s n)

pcoe_ok N.f : N2 → N2
-- `N.f ↑n` can be coerced as `⇡(N.f ↑n) : N2` for `n : N2`.
-- Consequently, `N.f (N.f ↑n)` can be coerced as `⇡(N.f ↑⇡(N.f ↑n))` for `n : N2`
-- (which simplifies to `⇡(N.f (N.f ↑n))`).

Also, some more examples:

def N.add : N → N → N
| n, z      => n
| n, (s n') => s (N.add n n')
pcoe_ok N.add : N2 → N2 → N2

def N.mul : N → N → N
  | _, z   => z
  | a, s b => N.add (N.mul a b) a
pcoe_ok N.mul : N2 → N2 → N2
pcoe_ok N.mul : N → N2 → N2
pcoe_ok N.mul : N2 → N → N2

/--
Safety wrapper around `N` for natural numbers greater than some `n`.
-/
[@safety N] -- adds the `Coe`/`PCoe` instances
structure NGT (n : N) : Type where
toN : N

pcoe_ok N.s : NGT → NGT
-- `N.s ↑n` can be coerced as `⇡(N.s ↑n) : NGT m` for `n : NGT m`.

/--
Safety wrapper around `N` for natural numbers less than some `n`.
-/
[@safety N]
structure NLT (n : N) : Type where
toN : N

pcoe_ok N.z : NLT
pcoe_ok_ctor N.s : NLT → NLT
-- within the `N.s` minor premise of `N.rec`, the reflexive instance of `N` can be cast:
def dec : NLT m → NLT m
| .z => .z -- elabs to (⇡z : NLT m)
| .s n => n -- elabs to (⇡n : NLT m)