anonymous structure

Stream: lean4

[Tomas Skrivan]

For quite some time I wanted to have anonymous structures.

My main use case is when a function returns multiple values. Often I end up defining one time use structure likeFooResult for function foo. However more often I just end up returning a type that looks like _ × _ × _ × _ × _ which I often regret when I read the code after some time.

So I finally took the time and defined anonymous structures. You can write code like

def foo (a : Float) : struct {x : Float, y : Float} :=
  { x := a.cos, y := a.sin }

instance : Add (struct { x : Float, y : Float }) := ⟨fun u v => ⟨u.x+v.x, u.y+v.y⟩⟩

/-- info: { x := 1.916065, y := 0.489252 } -/
#guard_msgs in
#eval
  let a := foo 0.1
  let b := foo 0.4
  a + b

No question here, I just want to share this as I'm sure someone else will find it useful too.

Here is the implementation

code

import Lean
import Std.Data.HashMap

open Lean Elab Command Term Meta Parser

namespace AnonymousStruct

-- Field information for the struct
structure FieldInfo where
  name : Name
  type : Expr
  deriving Inhabited, BEq, Hashable

instance : ToString FieldInfo where
  toString fi := s!"{fi.name} : {fi.type}"

-- Struct signature for lookup
structure StructSignature where
  fields : Array FieldInfo
  deriving Inhabited, BEq, Hashable

-- Entry in the environment extension
structure AnonStructEntry where
  signature : StructSignature
  structName : Name
  deriving Inhabited

-- Environment extension to track anonymous structs
initialize anonStructExt : SimplePersistentEnvExtension AnonStructEntry (Std.HashMap StructSignature Name) ←
  registerSimplePersistentEnvExtension {
    addImportedFn := fun as =>
      as.foldl (fun acc a => a.foldl (fun m e => m.insert e.signature e.structName) acc) {}
    addEntryFn := fun m e => m.insert e.signature e.structName
  }

-- Get or create a structure for the given signature
def getOrCreateStruct (fields : Array FieldInfo) : CommandElabM Name := do
  let sig : StructSignature := { fields := fields }
  -- let normSig := sig.normalize

  -- Check if we already have this structure
  let env ← getEnv
  let structMap := anonStructExt.getState env

  if let some existingName := structMap[sig]? then
    return existingName

  -- Generate a fresh name for the structure
  let mut freshName := `AnonStruct
  let mut idx := structMap.size
  while (← getEnv).contains freshName do
    idx := idx + 1
    freshName := (`AnonStruct).appendIndexAfter idx

  -- Build field declarations for the structure
  let mut fieldDecls : Array (TSyntax ``Command.structExplicitBinder) := #[]
  for fi in fields do
    let fieldName := mkIdent fi.name
    let fieldType := fi.type
    let fieldTypeStx ← liftTermElabM
      <| withOptions (fun opt => opt.setBool `pp.notation false)
      <| PrettyPrinter.delab fieldType
    fieldDecls := fieldDecls.push (← `(Command.structExplicitBinder| ($fieldName:ident : $fieldTypeStx)))

  -- Create and elaborate the structure command
  let structCmd ← `(command| structure $(mkIdent freshName):ident where
    $[$fieldDecls]*)

  elabCommand structCmd

  -- Register in our extension
  modifyEnv fun env =>
    anonStructExt.addEntry env { signature := sig, structName := freshName }

  return freshName

-- Syntax for anonymous structs
declare_syntax_cat anonym_struct_field
syntax ident ":" term : anonym_struct_field

/-- Anonymous structure. -/
syntax (name := anonStructType) "struct " "{" anonym_struct_field,* "}" : term

-- Macro to elaborate anonymous struct types
elab_rules : term
| `(struct { $[$names:ident : $types:term],* }) => do
  let types : Array Term := types
  let stx ← `(term| struct { $[$names:ident : $types:term],* })
  let types ← types.mapM (elabTerm · none)
  let fields := names.zip types |>.map fun (n, t) =>
    ({ name := n.getId, type := t } : FieldInfo)

  let structName ← liftCommandElabM (getOrCreateStruct fields)
  let structId := mkIdent structName
  let cmd ← `(command|
    @[app_unexpander $structId]
    def unexpandAnonStruct : Lean.PrettyPrinter.Unexpander
    | _ => `($stx)
    )
  let _ ← liftCommandElabM (elabCommand cmd)
  return (.const structName [])

end AnonymousStruct

[Kenny Lau]

very nice! just one tiny detail, I would write struct% to make it clear that this is a meta function

[Tomas Skrivan]

Ohh one question regardless. Can someone help me to fix the unexpander? It uses syntax quotation inside of quotation and all the identifiers get murdered with ✝

/-- info: struct {x✝:Float✝, y✝:Float✝} : Type -/
#check struct { x : Float, y : Float }

[Kenny Lau]

instead of `(struct { $[$names:ident : $types:term],* }) I would do the parsing separately, maybe this would fix that issue

[Tomas Skrivan]

Kenny Lau said:

very nice! just one tiny detail, I would write struct% to make it clear that this is a meta function

Yeah that would work too but I feel that the % notation is used for elaborators that do substantially more complicated things. I feel that anonymous struct is quite a simple concept.

[Tomas Skrivan]

Kenny Lau said:

instead of `(struct { $[$names:ident : $types:term],* }) I would do the parsing separately, maybe this would fix that issue

What do you mean by passing separately?

I might be wrong but I thought the issue is in

let cmd ← `(command|
    @[app_unexpander $structId]
    def unexpandAnonStruct : Lean.PrettyPrinter.Unexpander
    | _ => `($stx)
    )

which contains nested quotation `

[František Silváši 🦉]

I wonder if this isn't a legitimate use of suspending hygiene:

@[app_unexpander $structId]
    def unexpandAnonStruct : Lean.PrettyPrinter.Unexpander
    | _ => set_option hygiene false in `($stx)
    )

Considering the macro scopes have already been 'handled' by this point (i.e. stx is the thing you want verbatim here)?

[Tomas Skrivan]

Thanks a lot! works like a charm!

Yes it feels like a legitimate use case but I don't really understand syntax and macro scopes so hard to say.