Omit.lean

import Lean

open Lean Elab Command Meta

-- Macro to create a structure by omitting specified fields
syntax "elab_omit" ident "[" term,* "]" : term

-- Macro to create a structure by picking specified fields  
syntax "elab_pick" ident "[" term,* "]" : term

-- Helper to get structure fields
def getStructureFields (structName : Name) : MetaM (Array Name) := do
  let env ← getEnv
  match getStructureInfo? env structName with
  | some info => return info.fieldNames
  | none => throwError s!"'{structName}' is not a structure"

-- Helper to get field type
def getFieldType (structName : Name) (fieldName : Name) : MetaM Expr := do
  let env ← getEnv
  let projName := structName ++ fieldName
  match env.find? projName with
  | some (ConstantInfo.defnInfo val) => return val.type
  | _ => throwError s!"Could not find field type for {fieldName}"

-- Command to define structure with omitted fields
elab "structure " newName:ident " := " "elab_omit " baseName:ident " [" omitFields:term,* "]" : command => do
  let baseNameId := baseName.getId
  let newNameId := newName.getId
  
  -- Get all fields from base structure
  let allFields ← liftMetaM <| getStructureFields baseNameId
  
  -- Parse fields to omit
  let omitList := omitFields.getElems.map (·.raw.getId)
  
  -- Filter fields
  let keepFields := allFields.filter (fun f => !omitList.contains f)
  
  if keepFields.isEmpty then
    throwError "Cannot create structure with no fields"
  
  -- Build field declarations
  let mut fieldStrs := #[]
  for field in keepFields do
    let fullFieldName := baseNameId ++ field
    let env ← getEnv
    match env.find? fullFieldName with
    | some info => 
      let typeStr ← liftMetaM <| Meta.ppExpr info.type
      fieldStrs := fieldStrs.push s!"  {field} : {typeStr}"
    | none => throwError s!"Field {field} not found"
  
  -- Create structure command
  let structDef := String.intercalate "\n" fieldStrs.toList
  let cmd := s!"structure {newNameId} where\n{structDef}\nderiving Repr"
  
  elabCommand (← `(command| $cmd:str))

-- Command to define structure with picked fields
elab "structure " newName:ident " := " "elab_pick " baseName:ident " [" pickFields:term,* "]" : command => do
  let baseNameId := baseName.getId
  let newNameId := newName.getId
  
  -- Get all fields from base structure
  let allFields ← liftMetaM <| getStructureFields baseNameId
  
  -- Parse fields to pick
  let pickList := pickFields.getElems.map (·.raw.getId)
  
  -- Validate all picked fields exist
  for field in pickList do
    if !allFields.contains field then
      throwError s!"Field '{field}' does not exist in {baseNameId}"
  
  if pickList.isEmpty then
    throwError "Must pick at least one field"
  
  -- Build field declarations
  let mut fieldStrs := #[]
  for field in pickList do
    let fullFieldName := baseNameId ++ field
    let env ← getEnv
    match env.find? fullFieldName with
    | some info =>
      let typeStr ← liftMetaM <| Meta.ppExpr info.type
      fieldStrs := fieldStrs.push s!"  {field} : {typeStr}"
    | none => throwError s!"Field {field} not found"
  
  -- Create structure command
  let structDef := String.intercalate "\n" fieldStrs.toList
  let cmd := s!"structure {newNameId} where\n{structDef}\nderiving Repr"
  
  elabCommand (← `(command| $cmd:str))

-- Test: Define base structure
structure Person where
  name : String
  age : Nat
  email : String
  address : String
deriving Repr

-- Use elab_omit to create structure without address
structure PersonWithoutAddress := elab_omit Person ["address"]

-- Use elab_omit to create structure without email and address
structure PersonBasic := elab_omit Person ["email", "address"]

-- Use elab_pick to create structure with only name and age
structure PersonNameAge := elab_pick Person ["name", "age"]

-- Use elab_pick to create structure with only contact info
structure PersonContact := elab_pick Person ["email", "address"]

-- Helper conversion functions
def Person.toWithoutAddress (p : Person) : PersonWithoutAddress :=
  { name := p.name, age := p.age, email := p.email }

def Person.toBasic (p : Person) : PersonBasic :=
  { name := p.name, age := p.age }

def Person.toNameAge (p : Person) : PersonNameAge :=
  { name := p.name, age := p.age }

def Person.toContact (p : Person) : PersonContact :=
  { email := p.email, address := p.address }

-- Test the structures
def main : IO Unit := do
  let person : Person := {
    name := "Alice"
    age := 30
    email := "alice@example.com"
    address := "123 Main St"
  }
  
  IO.println "=== Original Person ==="
  IO.println (repr person)
  
  IO.println "\n=== PersonWithoutAddress (omit address) ==="
  IO.println (repr person.toWithoutAddress)
  
  IO.println "\n=== PersonBasic (omit email, address) ==="
  IO.println (repr person.toBasic)
  
  IO.println "\n=== PersonNameAge (pick name, age) ==="
  IO.println (repr person.toNameAge)
  
  IO.println "\n=== PersonContact (pick email, address) ==="
  IO.println (repr person.toContact)

#eval main