From 8b248e2d8c7e265ea4ba0ed491cdf52daa7514a3 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Thu, 3 Mar 2022 19:52:26 -0800 Subject: [PATCH] feat: elaborate `for h : x in xs do ...` notation --- src/Lean/Elab/Do.lean | 15 ++++++++--- src/Lean/Elab/Extra.lean | 52 +++++++++++++++++++++++++++---------- tests/lean/run/wfForIn.lean | 9 +++++++ 3 files changed, 59 insertions(+), 17 deletions(-) create mode 100644 tests/lean/run/wfForIn.lean diff --git a/src/Lean/Elab/Do.lean b/src/Lean/Elab/Do.lean index 09ac798acd..ff2670803f 100644 --- a/src/Lean/Elab/Do.lean +++ b/src/Lean/Elab/Do.lean @@ -1415,8 +1415,7 @@ mutual do $body) doSeqToCode (getDoSeqElems (getDoSeq auxDo) ++ doElems) else withRef doFor do - unless doForDecls[0][0].isNone do - throwErrorAt doForDecls[0][0] "the proof annotation here has not been implemented yet" + let h? := if doForDecls[0][0].isNone then none else some doForDecls[0][0][0] let x := doForDecls[0][1] withRef x <| checkNotShadowingMutable (← getPatternVarsEx x) let xs := doForDecls[0][3] @@ -1426,7 +1425,11 @@ mutual let uvarsTuple ← liftMacroM do mkTuple (← uvars.mapM mkIdentFromRef) if hasReturn forInBodyCodeBlock.code then let forInBody ← liftMacroM <| destructTuple uvars (← `(r)) forInBody - let forInTerm ← `(for_in% $(xs) (MProd.mk none $uvarsTuple) fun $x r => let r := r.2; $forInBody) + let forInTerm ← + if let some h := h? then + `(for_in'% $(xs) (MProd.mk none $uvarsTuple) fun $x $h r => let r := r.2; $forInBody) + else + `(for_in% $(xs) (MProd.mk none $uvarsTuple) fun $x r => let r := r.2; $forInBody) let auxDo ← `(do let r ← $forInTerm:term; $uvarsTuple:term := r.2; match r.1 with @@ -1435,7 +1438,11 @@ mutual doSeqToCode (getDoSeqElems (getDoSeq auxDo) ++ doElems) else let forInBody ← liftMacroM <| destructTuple uvars (← `(r)) forInBody - let forInTerm ← `(for_in% $(xs) $uvarsTuple fun $x r => $forInBody) + let forInTerm ← + if let some h := h? then + `(for_in'% $(xs) $uvarsTuple fun $x $h r => $forInBody) + else + `(for_in% $(xs) $uvarsTuple fun $x r => $forInBody) if doElems.isEmpty then let auxDo ← `(do let r ← $forInTerm:term; $uvarsTuple:term := r; diff --git a/src/Lean/Elab/Extra.lean b/src/Lean/Elab/Extra.lean index 275bc4516a..52e48f9a22 100644 --- a/src/Lean/Elab/Extra.lean +++ b/src/Lean/Elab/Extra.lean @@ -58,6 +58,17 @@ where @[builtinTermElab binrel_no_prop] def elabBinRelNoProp : TermElab := elabBinRelCore true +private def getMonadForIn (expectedType? : Option Expr) : TermElabM Expr := do + match expectedType? with + | none => throwError "invalid 'for_in%' notation, expected type is not available" + | some expectedType => + match (← isTypeApp? expectedType) with + | some (m, _) => return m + | none => throwError "invalid 'for_in%' notation, expected type is not of of the form `M α`{indentExpr expectedType}" + +private def throwForInFailure (forInInstance : Expr) : TermElabM Expr := + throwError "failed to synthesize instance for 'for_in%' notation{indentExpr forInInstance}" + @[builtinTermElab forInMacro] def elabForIn : TermElab := fun stx expectedType? => do match stx with | `(for_in% $col $init $body) => @@ -65,7 +76,7 @@ where | none => elabTerm (← `(let col := $col; for_in% col $init $body)) expectedType? | some colFVar => tryPostponeIfNoneOrMVar expectedType? - let m ← getMonad expectedType? + let m ← getMonadForIn expectedType? let colType ← inferType colFVar let elemType ← mkFreshExprMVar (mkSort (mkLevelSucc (← mkFreshLevelMVar))) let forInInstance ← @@ -78,19 +89,34 @@ where let ref ← getRef let forInFn ← mkConst ``forIn elabAppArgs forInFn #[] #[Arg.stx col, Arg.stx init, Arg.stx body] expectedType? (explicit := false) (ellipsis := false) - | LOption.undef => tryPostpone; throwFailure forInInstance - | LOption.none => throwFailure forInInstance + | LOption.undef => tryPostpone; throwForInFailure forInInstance + | LOption.none => throwForInFailure forInInstance + | _ => throwUnsupportedSyntax + +@[builtinTermElab forInMacro'] def elabForIn' : TermElab := fun stx expectedType? => do + match stx with + | `(for_in'% $col $init $body) => + match (← isLocalIdent? col) with + | none => elabTerm (← `(let col := $col; for_in'% col $init $body)) expectedType? + | some colFVar => + tryPostponeIfNoneOrMVar expectedType? + let m ← getMonadForIn expectedType? + let colType ← inferType colFVar + let elemType ← mkFreshExprMVar (mkSort (mkLevelSucc (← mkFreshLevelMVar))) + let memType ← mkFreshExprMVar (← mkAppM ``Membership #[elemType, colType]) + let forInInstance ← + try + mkAppM ``ForIn' #[m, colType, elemType, memType] + catch + ex => tryPostpone; throwError "failed to construct `ForIn'` instance for collection{indentExpr colType}\nand monad{indentExpr m}" + match (← trySynthInstance forInInstance) with + | LOption.some val => + let ref ← getRef + let forInFn ← mkConst ``forIn' + elabAppArgs forInFn #[] #[Arg.expr colFVar, Arg.stx init, Arg.stx body] expectedType? (explicit := false) (ellipsis := false) + | LOption.undef => tryPostpone; throwForInFailure forInInstance + | LOption.none => throwForInFailure forInInstance | _ => throwUnsupportedSyntax -where - getMonad (expectedType? : Option Expr) : TermElabM Expr := do - match expectedType? with - | none => throwError "invalid 'for_in%' notation, expected type is not available" - | some expectedType => - match (← isTypeApp? expectedType) with - | some (m, _) => return m - | none => throwError "invalid 'for_in%' notation, expected type is not of of the form `M α`{indentExpr expectedType}" - throwFailure (forInInstance : Expr) : TermElabM Expr := - throwError "failed to synthesize instance for 'for_in%' notation{indentExpr forInInstance}" namespace BinOp /- diff --git a/tests/lean/run/wfForIn.lean b/tests/lean/run/wfForIn.lean new file mode 100644 index 0000000000..054d6630cf --- /dev/null +++ b/tests/lean/run/wfForIn.lean @@ -0,0 +1,9 @@ +inductive Term where + | app (f : String) (args : List Term) + +def printFns : Term → IO Unit + | Term.app f args => do + IO.println f + for h : arg in args do + have : sizeOf arg < 1 + sizeOf f + sizeOf args := Nat.lt_trans (List.sizeOf_lt_of_mem h) (by simp_arith) + printFns arg