From 074259201e30bcd554c134ad5359ed7e97885b37 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Sat, 19 Dec 2020 14:09:53 -0800 Subject: [PATCH] feat: add helper classes for implementing parallel `for` It is based on an approach suggested by Andrew Kent, and refined by Sebastian Ullrich. TODO: expand the the parallel `for`s at `Do.lean`. --- src/Init/Data.lean | 1 + src/Init/Data/Stream.lean | 100 ++++++++++++++++++++++++++++ tests/lean/stream.lean | 40 +++++++++++ tests/lean/stream.lean.expected.out | 31 +++++++++ 4 files changed, 172 insertions(+) create mode 100644 src/Init/Data/Stream.lean create mode 100644 tests/lean/stream.lean create mode 100644 tests/lean/stream.lean.expected.out diff --git a/src/Init/Data.lean b/src/Init/Data.lean index 02fe948ce1..ecb12e592a 100644 --- a/src/Init/Data.lean +++ b/src/Init/Data.lean @@ -23,3 +23,4 @@ import Init.Data.Range import Init.Data.Hashable import Init.Data.OfScientific import Init.Data.Format +import Init.Data.Stream diff --git a/src/Init/Data/Stream.lean b/src/Init/Data/Stream.lean new file mode 100644 index 0000000000..97c874b386 --- /dev/null +++ b/src/Init/Data/Stream.lean @@ -0,0 +1,100 @@ +/- +Copyright (c) 2020 Microsoft Corporation. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Sebastian Ullrich, Andrew Kent, Leonardo de Moura +-/ +prelude +import Init.Data.Array.Subarray +import Init.Data.Range + +/- + Streams are used to implement parallel `for` statements. + Example: + ``` + for x in xs, y in ys do + ... + ``` + is expanded into + ``` + let mut s := toStream ys + for x in xs do + match Stream.next? s with + | none => break + | some (y, s') => + s := s' + ... + ``` +-/ +class ToStream (collection : Type u) (stream : outParam (Type u)) : Type u where + toStream : collection → stream + +export ToStream (toStream) + +class Stream (stream : Type u) (value : outParam (Type v)) : Type (max u v) where + next? : stream → Option (value × stream) + +/- Helper class for using dot-notation with `Stream`s -/ +structure StreamOf (ρ : Type u) where + s : ρ + +abbrev streamOf (s : ρ) := + StreamOf.mk s + +@[inline] partial def StreamOf.forIn [Stream ρ α] [Monad m] [Inhabited (m β)] (s : StreamOf ρ) (b : β) (f : α → β → m (ForInStep β)) : m β := do + let rec @[specialize] visit (s : ρ) (b : β) : m β := do + match Stream.next? s with + | some (a, s) => match (← f a b) with + | ForInStep.done b => return b + | ForInStep.yield b => visit s b + | none => return b + visit s.s b + +instance : ToStream (List α) (List α) where + toStream c := c + +instance : ToStream (Array α) (Subarray α) where + toStream a := a[:a.size] + +instance : ToStream (Subarray α) (Subarray α) where + toStream a := a + +instance : ToStream String Substring where + toStream s := s.toSubstring + +instance : ToStream Std.Range Std.Range where + toStream r := r + +instance [Stream ρ α] [Stream γ β] : Stream (ρ × γ) (α × β) where + next? | (s₁, s₂) => + match Stream.next? s₁ with + | none => none + | some (a, s₁) => match Stream.next? s₂ with + | none => none + | some (b, s₂) => some ((a, b), (s₁, s₂)) + +instance : Stream (List α) α where + next? + | [] => none + | a::as => some (a, as) + +instance : Stream (Subarray α) α where + next? s := + if h : s.start < s.stop then + have s.start + 1 ≤ s.stop from Nat.succLeOfLt h + some (s.as.get ⟨s.start, Nat.ltOfLtOfLe h s.h₂⟩, { s with start := s.start + 1, h₁ := this }) + else + none + +instance : Stream Std.Range Nat where + next? r := + if r.start < r.stop then + some (r.start, { r with start := r.start + r.step }) + else + none + +instance : Stream Substring Char where + next? s := + if s.startPos < s.stopPos then + some (s.str.get s.startPos, { s with startPos := s.str.next s.startPos }) + else + none diff --git a/tests/lean/stream.lean b/tests/lean/stream.lean new file mode 100644 index 0000000000..5a9dc09e9b --- /dev/null +++ b/tests/lean/stream.lean @@ -0,0 +1,40 @@ +def g (xs : Array Nat) (ys : Array Nat) : IO Unit := do + let mut s := toStream ys + for x in xs do + match Stream.next? s with + | none => break + | some (y, s') => + s := s' + IO.println s!"x: {x}, y: {y}" + +#eval g #[1, 2, 3] #[4, 5] +#print "-----" + +def f [Stream ρ α] [ToString α] (xs : Array Nat) (ys : ρ) : IO Unit := do + let mut ys := ys + for x in xs do + match Stream.next? ys with + | none => break + | some (y, ys') => + ys := ys' + IO.println s!"x: {x}, y: {y}" + +#eval f #[1, 2, 3] (toStream #[4, 5]) +#print "-----" +#eval f #[1, 2] (toStream [4, 5, 6]) +#print "-----" +#eval f #[1, 2, 3] (toStream "hello") + +def h [Stream ρ α] [ToString α] (s : ρ) : IO Unit := do + for a in streamOf s do + IO.println a + +#eval h (toStream [1, 2, 3]) +#print "-----" +#eval h [1, 2, 3] +#print "-----" +#eval h (toStream #[1, 2, 3]) +#print "-----" +#eval h #[1, 2, 3, 4, 5, 6, 7][2:5] +#print "-----" +#eval h ([1, 2, 3], [4, 5, 6]) diff --git a/tests/lean/stream.lean.expected.out b/tests/lean/stream.lean.expected.out new file mode 100644 index 0000000000..7926806095 --- /dev/null +++ b/tests/lean/stream.lean.expected.out @@ -0,0 +1,31 @@ +x: 1, y: 4 +x: 2, y: 5 +----- +x: 1, y: 4 +x: 2, y: 5 +----- +x: 1, y: 4 +x: 2, y: 5 +----- +x: 1, y: h +x: 2, y: e +x: 3, y: l +1 +2 +3 +----- +1 +2 +3 +----- +1 +2 +3 +----- +3 +4 +5 +----- +(1, 4) +(2, 5) +(3, 6)