nandi/rustnim · 5861a17

⑂ 5861a17

Lower `.take(n)` and a block used as a value

`.take(n)` is the first `n` elements, or all of them if there are fewer. That
is a shorter index loop over the same container, so it joins the chain like
`enumerate` and `zip` rather than materialising anything -- and the binding is
the inner chain's, since taking fewer elements changes nothing about what each
one is.

A block is an expression in Rust and Nim's statements are not, so the
statements run where they stand and the trailing expression is the value. The
single-expression case was already handled; this is the rest of it.

tests/cases/080.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

nandithebull committed 2026-09-20T14:33:17-07:00 Browse files

5861a17 parent: 7303426

modified src/lower.rs +32 -2 Viewed

@@ -118,6 +118,10 @@ enum Iter {
118	118	/// lazy iterator this is exact; it is only in an index loop that it
119	119	/// changes the loop's shape.
120	120	Filter { inner: Box<Iter>, param: String, body: Box<Expr> },
	121	+ /// `.take(n)` -- the first `n` elements, or all of them if there are
	122	+ /// fewer. That is a shorter index loop over the same container, so it
	123	+ /// composes with the rest rather than materialising anything.
	124	+ Take { inner: Box<Iter>, n: String },
121	125	}
122	126
123	127	impl Iter {
@@ -134,6 +138,7 @@ impl Iter {
134	138	Iter::Windows { len, k, .. } => format!("(max(0, {} - int({}) + 1))", len, k),
135	139	Iter::Enumerate(i) => i.len(),
136	140	Iter::Map { inner, .. } \| Iter::Filter { inner, .. } => inner.len(),
	141	+ Iter::Take { inner, n } => format!("min({}, int({}))", inner.len(), n),
137	142	// A closure iterator has no length; it is iterated directly.
138	143	Iter::Closure { .. } => "0".into(),
139	144	Iter::Zip(a, b) => format!("min({}, {})", a.len(), b.len()),
@@ -145,6 +150,7 @@ impl Iter {
145	150	match self {
146	151	Iter::Elems { elem, .. } => elem.clone(),
147	152	Iter::Enumerate(i) => i.elem(),
	153	+ Iter::Take { inner, .. } => inner.elem(),
148	154	Iter::Closure { elem, .. } => elem.clone(),
149	155	_ => None,
150	156	}
@@ -155,6 +161,7 @@ impl Iter {
155	161	match self {
156	162	Iter::Elems { code, off, .. } => format!("{}[int({}) + {}]", code, off, i),
157	163	Iter::Enumerate(inner) => inner.at(i),
	164	+ Iter::Take { inner, .. } => inner.at(i),
158	165	_ => i.to_string(),
159	166	}
160	167	}
@@ -5459,6 +5466,13 @@ impl Lowerer {
5459	5466	mutable: false,
5460	5467	})
5461	5468	}
	5469	+ "take" if m.args.len() == 1 => {
	5470	+ let n = self.expr(&m.args[0])?;
	5471	+ Ok(Iter::Take {
	5472	+ inner: Box::new(self.resolve_iter(&m.receiver)?),
	5473	+ n: n.code,
	5474	+ })
	5475	+ }
5462	5476	"enumerate" if m.args.is_empty() => {
5463	5477	Ok(Iter::Enumerate(Box::new(self.resolve_iter(&m.receiver)?)))
5464	5478	}
@@ -5806,6 +5820,9 @@ impl Lowerer {
5806	5820	// Nim is already the value.
5807	5821	(Pat::Reference(r), _) => self.bind_pattern(&r.pat, it, i),
5808	5822	(Pat::Paren(p), _) => self.bind_pattern(&p.pat, it, i),
	5823	+ // `.take(n)` shortens the loop and changes nothing about what each
	5824	+ // element is, so the binding is the inner chain's.
	5825	+ (p, Iter::Take { inner, .. }) => self.bind_pattern(p, inner, i),
5809	5826	(Pat::Ident(id), _) => {
5810	5827	let name = id.ident.to_string();
5811	5828	match it {
@@ -5877,7 +5894,8 @@ impl Lowerer {
5877	5894	}
5878	5895	// Handled above.
5879	5896	Iter::Zip(..) \| Iter::Enumerate(..) \| Iter::Map { .. }
5880		- \| Iter::Filter { .. } \| Iter::Closure { .. } => unreachable!(),
	5897	+ \| Iter::Filter { .. } \| Iter::Closure { .. }
	5898	+ \| Iter::Take { .. } => unreachable!(),
5881	5899	}
5882	5900	}
5883	5901	_ => Err("unsupported `for` pattern".into()),
@@ -6910,9 +6928,14 @@ impl Lowerer {
6910	6928	if let Some(Stmt::Expr(e, None)) = b.block.stmts.first() {
6911	6929	self.expr_at(e, expect)
6912	6930	} else {
6913		- Err("block expression with statements in value position is not implemented yet".into())
	6931	+ self.block_value(&b.block, expect)
6914	6932	}
6915	6933	}
	6934	+ // A block used as a value: the statements run where they stand and
	6935	+ // the trailing expression is the value. Rust's block is an
	6936	+ // expression; Nim's statements are not, so the two are separated
	6937	+ // rather than nested.
	6938	+ Expr::Block(b) => self.block_value(&b.block, expect),
6916	6939	// An assignment is an expression in Rust, of type `()`. In Nim it
6917	6940	// is a statement, so it is emitted as one and the value it stands
6918	6941	// for is the unit -- which is what Rust's is.
@@ -7683,6 +7706,13 @@ impl Lowerer {
7683	7706	))
7684	7707	}
7685	7708
	7709	+ fn block_value(&mut self, b: &syn::Block, expect: Option<&Nim>) -> Result<Val, String> {
	7710	+ self.push_scope();
	7711	+ let r = self.block_body_at(b, expect);
	7712	+ self.pop_scope();
	7713	+ r?.ok_or_else(\|\| "a block used as a value has to end in one".to_string())
	7714	+ }
	7715	+
7686	7716	/// Lower a block's statements at the current indentation, without opening
7687	7717	/// a Nim `block:` -- used for `unsafe { .. }`, which introduces no scope
7688	7718	/// of its own in the generated code.

      /// lazy iterator this is exact; it is only in an index loop that it
      /// changes the loop's shape.
      Filter { inner: Box<Iter>, param: String, body: Box<Expr> },
++    /// `.take(n)` -- the first `n` elements, or all of them if there are
++    /// fewer. That is a shorter index loop over the same container, so it
++    /// composes with the rest rather than materialising anything.
++    Take { inner: Box<Iter>, n: String },
  }
  impl Iter {
              Iter::Windows { len, k, .. } => format!("(max(0, {} - int({}) + 1))", len, k),
              Iter::Enumerate(i) => i.len(),
              Iter::Map { inner, .. } | Iter::Filter { inner, .. } => inner.len(),
++            Iter::Take { inner, n } => format!("min({}, int({}))", inner.len(), n),
              // A closure iterator has no length; it is iterated directly.
              Iter::Closure { .. } => "0".into(),
              Iter::Zip(a, b) => format!("min({}, {})", a.len(), b.len()),
          match self {
              Iter::Elems { elem, .. } => elem.clone(),
              Iter::Enumerate(i) => i.elem(),
++            Iter::Take { inner, .. } => inner.elem(),
              Iter::Closure { elem, .. } => elem.clone(),
              _ => None,
          }
          match self {
              Iter::Elems { code, off, .. } => format!("{}[int({}) + {}]", code, off, i),
              Iter::Enumerate(inner) => inner.at(i),
++            Iter::Take { inner, .. } => inner.at(i),
              _ => i.to_string(),
          }
      }
                              mutable: false,
                          })
                      }
++                    "take" if m.args.len() == 1 => {
++                        let n = self.expr(&m.args[0])?;
++                        Ok(Iter::Take {
++                            inner: Box::new(self.resolve_iter(&m.receiver)?),
++                            n: n.code,
++                        })
++                    }
                      "enumerate" if m.args.is_empty() => {
                          Ok(Iter::Enumerate(Box::new(self.resolve_iter(&m.receiver)?)))
                      }
              // Nim is already the value.
              (Pat::Reference(r), _) => self.bind_pattern(&r.pat, it, i),
              (Pat::Paren(p), _) => self.bind_pattern(&p.pat, it, i),
++            // `.take(n)` shortens the loop and changes nothing about what each
++            // element is, so the binding is the inner chain's.
++            (p, Iter::Take { inner, .. }) => self.bind_pattern(p, inner, i),
              (Pat::Ident(id), _) => {
                  let name = id.ident.to_string();
                  match it {
                      }
                      // Handled above.
                      Iter::Zip(..) | Iter::Enumerate(..) | Iter::Map { .. }
--                    | Iter::Filter { .. } | Iter::Closure { .. } => unreachable!(),
++                    | Iter::Filter { .. } | Iter::Closure { .. }
++                    | Iter::Take { .. } => unreachable!(),
                  }
              }
              _ => Err("unsupported `for` pattern".into()),
                  if let Some(Stmt::Expr(e, None)) = b.block.stmts.first() {
                      self.expr_at(e, expect)
                  } else {
--                    Err("block expression with statements in value position is not implemented yet".into())
++                    self.block_value(&b.block, expect)
                  }
              }
++            // A block used as a value: the statements run where they stand and
++            // the trailing expression is the value. Rust's block is an
++            // expression; Nim's statements are not, so the two are separated
++            // rather than nested.
++            Expr::Block(b) => self.block_value(&b.block, expect),
              // An assignment is an expression in Rust, of type `()`. In Nim it
              // is a statement, so it is emitted as one and the value it stands
              // for is the unit -- which is what Rust's is.
          ))
      }
++    fn block_value(&mut self, b: &syn::Block, expect: Option<&Nim>) -> Result<Val, String> {
++        self.push_scope();
++        let r = self.block_body_at(b, expect);
++        self.pop_scope();
++        r?.ok_or_else(|| "a block used as a value has to end in one".to_string())
++    }
++
      /// Lower a block's statements at the current indentation, without opening
      /// a Nim `block:` -- used for `unsafe { .. }`, which introduces no scope
      /// of its own in the generated code.

added tests/cases/080-take.rs +49 -0 Viewed

new file mode 100644
@@ -0,0 +1,49 @@
	1	+// `.take(n)` is the first `n` elements, or all of them if there are fewer —
	2	+// a shorter index loop over the same container, so it composes with the rest
	3	+// rather than materialising anything.
	4	+// A block used as a value: the statements run where they stand and the
	5	+// trailing expression is the value.
	6	+fn stepped(n: u32) -> u32 {
	7	+ let doubled = {
	8	+ let half = n / 2;
	9	+ let rest = n - half;
	10	+ half + rest * 2
	11	+ };
	12	+ doubled + 1
	13	+}
	14	+
	15	+fn main() {
	16	+ println!("{} {}", stepped(10), stepped(7));
	17	+
	18	+ let v: Vec<u32> = vec![10, 20, 30, 40, 50];
	19	+
	20	+ let mut t: u32 = 0;
	21	+ for x in v.iter().take(3) {
	22	+ t += *x;
	23	+ }
	24	+ println!("{}", t);
	25	+
	26	+ // More than there are.
	27	+ let mut u: u32 = 0;
	28	+ for x in v.iter().take(99) {
	29	+ u += *x;
	30	+ }
	31	+ println!("{}", u);
	32	+
	33	+ // None at all.
	34	+ let mut z: u32 = 0;
	35	+ for x in v.iter().take(0) {
	36	+ z += *x;
	37	+ }
	38	+ println!("{}", z);
	39	+
	40	+ // Composed with `enumerate`.
	41	+ for (i, x) in v.iter().take(2).enumerate() {
	42	+ println!("{} {}", i, x);
	43	+ }
	44	+
	45	+ for i in (0..10).take(4) {
	46	+ print!("{} ", i);
	47	+ }
	48	+ println!();
	49	+}

++// `.take(n)` is the first `n` elements, or all of them if there are fewer —
++// a shorter index loop over the same container, so it composes with the rest
++// rather than materialising anything.
++// A block used as a value: the statements run where they stand and the
++// trailing expression is the value.
++fn stepped(n: u32) -> u32 {
++    let doubled = {
++        let half = n / 2;
++        let rest = n - half;
++        half + rest * 2
++    };
++    doubled + 1
++}
++
++fn main() {
++    println!("{} {}", stepped(10), stepped(7));
++
++    let v: Vec<u32> = vec![10, 20, 30, 40, 50];
++
++    let mut t: u32 = 0;
++    for x in v.iter().take(3) {
++        t += *x;
++    }
++    println!("{}", t);
++
++    // More than there are.
++    let mut u: u32 = 0;
++    for x in v.iter().take(99) {
++        u += *x;
++    }
++    println!("{}", u);
++
++    // None at all.
++    let mut z: u32 = 0;
++    for x in v.iter().take(0) {
++        z += *x;
++    }
++    println!("{}", z);
++
++    // Composed with `enumerate`.
++    for (i, x) in v.iter().take(2).enumerate() {
++        println!("{} {}", i, x);
++    }
++
++    for i in (0..10).take(4) {
++        print!("{} ", i);
++    }
++    println!();
++}