Different elaboration inside/outside of match patterns

Stream: lean4

[Son Ho]

I defined a simple U32 machine integer type below, with some synctactic sugar to easily define constants:

import Lean

def U32 := { v : Int // 0 ≤ v ∧ v ≤ 4294967295 }

-- Notation for constants
notation:70 x:70 "#u32" => ⟨ x, by decide ⟩

example : U32 := 3#u32

-- We have to introduce a different notation for match patterns
notation:70 x:70 "#pat" => Subtype.mk x _

def is_zero (x : U32) : Bool :=
  match x with
  | 0#pat => true
  | _ => false

#eval is_zero (0#u32) -- true
#eval is_zero (1#u32) -- false

I need to introduce two notations: #u32 for the "regular" constants, and #pat for the case where we are inside a pattern. Is there a way of having a single notation, but which gets elaborated to different terms depending on whether we are inside a match pattern or not?

[Jannis Limperg]

Should be doable, but it's not entirely trivial. You'd have to define a custom elaborator (elab). This operates in the docs#Lean.Elab.Term.TermElabM monad, which gives you the info whether you're in a pattern or not. I can take a stab at this later.

[Son Ho]

That would be great, thanks! :)

[Son Ho]

But I should be able to manage provided I have a few pointers ;)

[Kyle Miller]

By the way, Lean 4 comes with docs#UInt32, and it has special runtime support.

[Kyle Miller]

You don't really need #u32, if you're ok with wraparound on overflow.

In patterns, you can just write 0 rather than 0#pat.

-- Notation for constants
notation:70 x:70 "#u32" => (⟨ x, by decide ⟩ : UInt32)

example : UInt32 := 3#u32

def is_zero (x : UInt32) : Bool :=
  match x with
  | 0 => true
  | _ => false

#eval is_zero (0#u32) -- true
#eval is_zero (1#u32) -- false

[Son Ho]

Kyle Miller said:

By the way, Lean 4 comes with docs#UInt32, and it has special runtime support.

Yes, I'm aware of this, but it doesn't fit my use case.
The example I show above is adapted from the Aeneas project whose goal is to generate models of Rust programs for the purpose of verification. In Rust, overflows lead to panics, so we do need to model them by not wraping around (and using an error monad).

[Mario Carneiro]

you should do this either by writing functions on UInt32 with the desired (non-)wrapping behavior, or by newtype-wrapping UInt32 so you can define what + etc do

[Son Ho]

Yes, but the example above is minimized: in practice I have to consider a whole family of machine integers, which include non-signed machine integers like i8 or i128, meaning I can't simply reuse a type like UInt32 or even Fin (the real definition I manipulate is parameterized by the bounds).

[Jannis Limperg]

Jannis Limperg said:

Should be doable, but it's not entirely trivial. You'd have to define a custom elaborator (elab). This operates in the docs#Lean.Elab.Term.TermElabM monad, which gives you the info whether you're in a pattern or not. I can take a stab at this later.

Doesn't look like this is possible actually. Lean analyses pattern Syntax here and only allows certain Syntax in patterns. So macros work for patterns, but custom elabs generally don't. Here's how far I got:

import Lean

open Lean Lean.Elab Lean.Elab.Term

def U32 := { v : Int // 0 ≤ v ∧ v ≤ 4294967295 }

elab x:num noWs "#u32" : term => do
  let stx ← if (← read).inPattern then
    `(Subtype.mk $x _)
  else
    `(⟨$x, by decide⟩)
  trace[debug] "before elabTerm"
  let e ← elabTerm stx $ some $ .const ``U32 []
  trace[debug] "after elabTerm"
  return e

def test : U32 := 3#u32

set_option trace.Elab.match true in
set_option trace.debug true in
def is_zero (x : U32) : Bool :=
  match x with
  | 0#u32 => true
  | _ => false

#eval is_zero (0#u32) -- true
#eval is_zero (1#u32) -- false

[Son Ho]

Thanks for trying!
I had a look at the collect function you pointed to and the fact that _ terms are treated differently in patterns gave me the idea to use typeclasses. Reusing this, I came up with the following (we solve the "in bounds" obligation by using a typeclass and Decidable: if outside a pattern Lean solves the typeclass, otherwise it treats it as a _ pattern):

open Lean Lean.Elab Lean.Elab.Term

def U32 := { v : Int // 0 ≤ v ∧ v ≤ 4294967295 }

class InBounds (v : Int) :=
  h : 0 ≤ v ∧ v ≤ 4294967295

@[match_pattern] def mkU32 (v : Int) [InBounds v] : U32 :=
  ⟨ v, InBounds.h ⟩

class DecideInBounds (p : Prop) [Decidable p] : Prop where
  isTrue : p

instance : @DecideInBounds p (.isTrue h) := @DecideInBounds.mk p (_) h

instance [DecideInBounds (0 ≤ v ∧ v ≤ 4294967295)] : InBounds v where
  h := DecideInBounds.isTrue

notation:max x:70 "#u32" => mkU32 x

def is_zero (x : U32) : Bool :=
  match x with
  | 0#u32 => true
  | _ => false

#eval is_zero 0#u32 -- true
#eval is_zero 1#u32 -- false

[Jannis Limperg]

Very clever. I just played around a bit with the code and discovered that it breaks when InBound is a Prop rather than (implicitly) a Type. :thinking:

[Son Ho]

Ok, that's interesting :sweat_smile: