feat: elaborate for h : x in xs do ... notation

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;

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
 /-

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