Make iter a builtin macro

Author: oli-obk

compiler/rustc_ast_lowering/src/expr.rs

@@ -1,4 +1,3 @@
-use std::assert_matches::assert_matches;
 use std::ops::ControlFlow;
 use std::sync::Arc;
 
@@ -1199,11 +1198,13 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let closure_def_id = self.local_def_id(closure_id);
         let (binder_clause, generic_params) = self.lower_closure_binder(binder);
 
-        assert_matches!(
-            coroutine_kind,
-            CoroutineKind::Async { .. },
-            "only async closures are supported currently"
-        );
+        let coroutine_desugaring = match coroutine_kind {
+            CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
+            CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
+            CoroutineKind::AsyncGen { span, .. } => {
+                span_bug!(span, "only async closures and `iter!` closures are supported currently")
+            }
+        };
 
         let body = self.with_new_scopes(fn_decl_span, |this| {
             let inner_decl =
@@ -1247,7 +1248,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             // Lower this as a `CoroutineClosure`. That will ensure that HIR typeck
             // knows that a `FnDecl` output type like `-> &str` actually means
             // "coroutine that returns &str", rather than directly returning a `&str`.
-            kind: hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Async),
+            kind: hir::ClosureKind::CoroutineClosure(coroutine_desugaring),
             constness: hir::Constness::NotConst,
         });
         hir::ExprKind::Closure(c)

compiler/rustc_borrowck/src/type_check/input_output.rs

@@ -52,7 +52,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             assert_matches!(
                 self.tcx().coroutine_kind(self.tcx().coroutine_for_closure(mir_def_id)),
                 Some(hir::CoroutineKind::Desugared(
-                    hir::CoroutineDesugaring::Async,
+                    hir::CoroutineDesugaring::Async | hir::CoroutineDesugaring::Gen,
                     hir::CoroutineSource::Closure
                 )),
                 "this needs to be modified if we're lowering non-async closures"

compiler/rustc_builtin_macros/src/iter.rs

@@ -0,0 +1,78 @@
+use rustc_ast::ptr::P;
+use rustc_ast::tokenstream::TokenStream;
+use rustc_ast::{
+    CaptureBy, ClosureBinder, Const, CoroutineKind, DUMMY_NODE_ID, Expr, ExprKind, ast, token,
+};
+use rustc_errors::PResult;
+use rustc_expand::base::{self, DummyResult, ExpandResult, ExtCtxt, MacroExpanderResult};
+use rustc_span::Span;
+
+pub(crate) fn expand<'cx>(
+    cx: &'cx mut ExtCtxt<'_>,
+    sp: Span,
+    tts: TokenStream,
+) -> MacroExpanderResult<'cx> {
+    let closure = match parse_closure(cx, sp, tts) {
+        Ok(parsed) => parsed,
+        Err(err) => {
+            return ExpandResult::Ready(DummyResult::any(sp, err.emit()));
+        }
+    };
+
+    ExpandResult::Ready(base::MacEager::expr(closure))
+}
+
+fn parse_closure<'a>(
+    cx: &mut ExtCtxt<'a>,
+    span: Span,
+    stream: TokenStream,
+) -> PResult<'a, P<Expr>> {
+    let mut parser = cx.new_parser_from_tts(stream);
+    let mut closure_parser = parser.clone();
+
+    let coroutine_kind = Some(CoroutineKind::Gen {
+        span,
+        closure_id: DUMMY_NODE_ID,
+        return_impl_trait_id: DUMMY_NODE_ID,
+    });
+
+    match closure_parser.parse_expr() {
+        Ok(mut closure) => {
+            if let ast::ExprKind::Closure(c) = &mut closure.kind {
+                if let Some(kind) = c.coroutine_kind {
+                    cx.dcx().span_err(kind.span(), "only plain closures allowed in `iter!`");
+                }
+                c.coroutine_kind = coroutine_kind;
+                if closure_parser.token != token::Eof {
+                    closure_parser.unexpected()?;
+                }
+                return Ok(closure);
+            }
+        }
+        Err(diag) => diag.cancel(),
+    }
+
+    let lo = parser.token.span.shrink_to_lo();
+    let block = parser.parse_block_tail(
+        lo,
+        ast::BlockCheckMode::Default,
+        rustc_parse::parser::AttemptLocalParseRecovery::No,
+    )?;
+    let fn_decl = cx.fn_decl(Default::default(), ast::FnRetTy::Default(span));
+    let closure = ast::Closure {
+        binder: ClosureBinder::NotPresent,
+        capture_clause: CaptureBy::Ref,
+        constness: Const::No,
+        coroutine_kind,
+        movability: ast::Movability::Movable,
+        fn_decl,
+        body: cx.expr_block(block),
+        fn_decl_span: span,
+        fn_arg_span: span,
+    };
+    if parser.token != token::Eof {
+        parser.unexpected()?;
+    }
+    let span = lo.to(parser.token.span);
+    Ok(cx.expr(span, ExprKind::Closure(Box::new(closure))))
+}

compiler/rustc_builtin_macros/src/lib.rs

@@ -49,6 +49,7 @@ mod errors;
 mod format;
 mod format_foreign;
 mod global_allocator;
+mod iter;
 mod log_syntax;
 mod pattern_type;
 mod source_util;
@@ -97,6 +98,7 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         include: source_util::expand_include,
         include_bytes: source_util::expand_include_bytes,
         include_str: source_util::expand_include_str,
+        iter: iter::expand,
         line: source_util::expand_line,
         log_syntax: log_syntax::expand_log_syntax,
         module_path: source_util::expand_mod,

compiler/rustc_hir_typeck/src/closure.rs

@@ -204,14 +204,19 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 )
             }
             hir::ClosureKind::CoroutineClosure(kind) => {
-                // async closures always return the type ascribed after the `->` (if present),
-                // and yield `()`.
                 let (bound_return_ty, bound_yield_ty) = match kind {
+                    hir::CoroutineDesugaring::Gen => {
+                        // `iter!` closures always return unit and yield the `Iterator::Item` type
+                        // that we have to infer.
+                        (tcx.types.unit, self.infcx.next_ty_var(expr_span))
+                    }
                     hir::CoroutineDesugaring::Async => {
+                        // async closures always return the type ascribed after the `->` (if present),
+                        // and yield `()`.
                         (bound_sig.skip_binder().output(), tcx.types.unit)
                     }
-                    hir::CoroutineDesugaring::Gen | hir::CoroutineDesugaring::AsyncGen => {
-                        todo!("`gen` and `async gen` closures not supported yet")
+                    hir::CoroutineDesugaring::AsyncGen => {
+                        todo!("`async gen` closures not supported yet")
                     }
                 };
                 // Compute all of the variables that will be used to populate the coroutine.

compiler/rustc_parse/src/parser/stmt.rs

@@ -714,7 +714,7 @@ impl<'a> Parser<'a> {
 
     /// Parses the rest of a block expression or function body.
     /// Precondition: already parsed the '{'.
-    pub(crate) fn parse_block_tail(
+    pub fn parse_block_tail(
         &mut self,
         lo: Span,
         s: BlockCheckMode,

library/core/src/iter/sources/generator.rs

@@ -20,6 +20,7 @@
 /// ```
 #[unstable(feature = "iter_macro", issue = "none", reason = "generators are unstable")]
 #[allow_internal_unstable(coroutines, iter_from_coroutine)]
+#[cfg_attr(not(bootstrap), rustc_builtin_macro)]
 pub macro iter($($t:tt)*) {
-    || $crate::iter::from_coroutine(#[coroutine] || { $($t)* })
+    /* compiler-builtin */
 }

tests/ui/iterators/generator.rs

@@ -21,8 +21,6 @@ fn main() {
     assert_eq!(i.next(), Some(16));
     assert_eq!(i.next(), Some(18));
     assert_eq!(i.next(), None);
-    // FIXME(iter_macro): desugar to `gen` instead of coroutines,
-    // as the latter panic after resuming iteration.
-    //assert_eq!(i.next(), None);
-    //assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
 }

tests/ui/iterators/generator_args.rs

@@ -0,0 +1,26 @@
+//@ run-pass
+
+#![feature(iter_macro)]
+// FIXME(iter_macro): make `yield` within it legal
+#![feature(coroutines)]
+
+use std::iter::iter;
+
+fn main() {
+    let i = iter! {|foo| {
+        yield foo;
+        for x in 5..10 {
+            yield x * 2;
+        }
+    }};
+    let mut i = i(3);
+    assert_eq!(i.next(), Some(3));
+    assert_eq!(i.next(), Some(10));
+    assert_eq!(i.next(), Some(12));
+    assert_eq!(i.next(), Some(14));
+    assert_eq!(i.next(), Some(16));
+    assert_eq!(i.next(), Some(18));
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+}

tests/ui/iterators/generator_capture.rs

@@ -0,0 +1,29 @@
+//@ run-pass
+
+#![feature(iter_macro)]
+// FIXME(iter_macro): make `yield` within it legal
+#![feature(coroutines)]
+
+use std::iter::iter;
+
+fn main() {
+    let i = {
+        let s = String::new();
+        iter! { move || {
+            yield s.len();
+            for x in 5..10 {
+                yield x * 2;
+            }
+        }}
+    };
+    let mut i = i();
+    assert_eq!(i.next(), Some(0));
+    assert_eq!(i.next(), Some(10));
+    assert_eq!(i.next(), Some(12));
+    assert_eq!(i.next(), Some(14));
+    assert_eq!(i.next(), Some(16));
+    assert_eq!(i.next(), Some(18));
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+}