Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Lazy infinite streams in Haskell, Swift, and Python
module Main where
import Prelude hiding (
iterate, repeat, map, filter, zipWith, zip,
take, takeWhile, drop, dropWhile)
import qualified Prelude as P
import Text.Printf (printf)
import Data.List (intercalate)
infixr 5 :>
data Stream a = a :> (Stream a)
first :: Stream a -> a
first (x :> _) = x
rest :: Stream a -> Stream a
rest (_ :> xs) = xs
iterate :: (a -> a) -> a -> Stream a
iterate f a = a :> iterate f (f a)
count :: Integer -> Integer -> Stream Integer
count start step = iterate (+step) start
repeat :: a -> Stream a
repeat a = a :> repeat a
map :: (a -> b) -> Stream a -> Stream b
map f (x :> xs) = f x :> map f xs
filter :: (a -> Bool) -> Stream a -> Stream a
filter p (x :> xs)
| p x = x :> filter p xs
| otherwise = filter p xs
zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith f (x :> xs) (y :> ys) = f x y :> zipWith f xs ys
zip :: Stream a -> Stream b -> Stream (a, b)
zip = zipWith (,)
scan :: (b -> a -> b) -> b -> Stream a -> Stream b
scan f z (x :> xs) = z :> scan f (f z x) xs
scan1 :: (a -> a -> a) -> Stream a -> Stream a
scan1 f (x :> xs) = scan f x xs
take :: Integer -> Stream a -> [a]
take n (x :> xs)
| n <= 0 = []
| otherwise = x : take (n - 1) xs
drop :: Integer -> Stream a -> Stream a
drop n xs
| n <= 0 = xs
| otherwise = drop (n - 1) (rest xs)
takeWhile :: (a -> Bool) -> Stream a -> [a]
takeWhile p (x :> xs)
| p x = x : takeWhile p xs
| otherwise = []
dropWhile :: (a -> Bool) -> Stream a -> Stream a
dropWhile p (x :> xs)
| p x = dropWhile p xs
| otherwise = x :> xs
annotate :: Show a => String -> [a] -> IO ()
annotate description body = putStrLn message
where
message = printf "%s:\n %s\n" description formatted
formatted = "[" ++ intercalate ", " (P.map show body) ++ "]"
nats = iterate (+1) 0
ones = repeat 1
strings = map (flip replicate '*') nats
squares = map (^2) nats
evens = filter even nats
odds = filter odd nats
sums = zipWith (+) nats nats
pairs = zip evens odds
facts = scan1 (*) (count 1 1)
fibs = 0 :> 1 :> zipWith (+) fibs (drop 1 fibs)
main = do
annotate
"First ten nats (iterate)"
(take 10 nats)
annotate
"First ten nats after dropping ten nats (drop)"
(take 10 (drop 10 nats))
annotate
"First ten ones (repeat)"
(take 10 ones)
annotate
"Type-changing map (map int -> string)"
(take 10 strings)
annotate
"First ten squares (map int -> int)"
(take 10 squares)
annotate
"First ten evens (filter int -> bool)"
(take 10 evens)
annotate
"First ten odds (filter int -> bool)"
(take 10 odds)
annotate
"First ten sums ((int, int) -> int)"
(take 10 sums)
annotate
"Even/odd pairs (zip)"
(take 10 pairs)
annotate
"First ten factorial (scan)"
(take 10 facts)
annotate
"First ten fibs (forward/tie/zipwith)"
(take 10 fibs)
annotate
"Take while <10"
(takeWhile (<10) nats)
annotate
"First ten after drop while <10"
(take 10 (dropWhile (<10) nats))
from __future__ import unicode_literals
import operator
class Stream(object):
sentinel = object()
__slots__ = ('_first', '_rest', '_step')
def __init__(self, first, rest):
self._first = first
if callable(rest):
self._step = rest
self._rest = self.sentinel
elif isinstance(rest, Stream):
self._step = lambda: rest
self._rest = rest
else:
raise TypeError('Second argument to Stream must be a callable or a Stream')
@property
def first(self):
return self._first
@property
def rest(self):
if self._rest is self.sentinel:
self._rest = self._step()
return self._rest
@classmethod
def cons(cls, head, rest):
return cls(head, rest)
@classmethod
def iterate(cls, f, x):
return cls(x, lambda: cls.iterate(f, f(x)))
@classmethod
def count(cls, start=0, step=1):
return cls.iterate(lambda i: i + step, start)
@classmethod
def repeat(cls, x):
return cls(x, lambda: cls.repeat(x))
def map(self, f):
return Stream(f(self.first), lambda: self.rest.map(f))
def filter(self, p):
if p(self.first):
return Stream(self.first, lambda: self.rest.filter(p))
return self.rest.filter(p)
def zipWith(self, f, stream):
return Stream(
f(self.first, stream.first),
lambda: self.rest.zipWith(f, stream.rest)
)
def zip(self, stream):
return self.zipWith(lambda x, y: (x, y), stream)
def scan(self, f, first):
return Stream(first, lambda: self.rest.scan(f, f(first, self.first)))
def scan1(self, f):
return self.rest.scan(f, self.first)
def take(self, n):
out = []
stream = self
i = 0
while i < n:
out.append(stream.first)
stream = stream.rest
i += 1
return out
def drop(self, n):
stream = self
i = 0
while i < n:
stream = stream.rest
i += 1
return stream
def takeWhile(self, p):
out = []
stream = self
while 1:
if p(stream.first):
out.append(stream.first)
stream = stream.rest
else:
break
return out
def dropWhile(self, p):
stream = self
while 1:
if p(stream.first):
stream = stream.rest
else:
break
return stream
def annotate(description, body):
formatted = ', '.join(map(str, body)).join('[]')
print("{}:\n {}\n".format(description, formatted))
nats = Stream.iterate(lambda x: x + 1, 0)
ones = Stream.repeat(1)
strings = nats.map(lambda x: '*' * x)
squares = nats.map(lambda x: x * x)
evens = nats.filter(lambda x: x % 2 == 0)
odds = nats.filter(lambda x: x % 2 == 1)
sums = nats.zipWith(operator.add, nats)
pairs = evens.zip(odds)
facts = Stream.count(1).scan1(operator.mul)
fibs = Stream.cons(0, Stream.cons(1, lambda: fibs.zipWith(operator.add, fibs.drop(1))))
if __name__ == '__main__':
annotate(
"First ten nats (iterate)",
nats.take(10),
)
annotate(
"First ten nats after dropping ten nats (drop)",
nats.drop(10).take(10),
)
annotate(
"First ten ones (repeat)",
ones.take(10),
)
annotate(
"Type-changing map (map int -> string)",
strings.take(10),
)
annotate(
"First ten squares (map int -> int)",
squares.take(10),
)
annotate(
"First ten evens (filter int -> bool)",
evens.take(10),
)
annotate(
"First ten odds (filter int -> bool)",
odds.take(10),
)
annotate(
"First ten sums ((int, int) -> int)",
sums.take(10),
)
annotate(
"Even/odd pairs (zip)",
pairs.take(10),
)
annotate(
"First ten factorial (scan)",
facts.take(10),
)
annotate(
"First ten fibs (forward/tie/zipwith)",
fibs.take(10),
)
annotate(
"Take while <10",
nats.takeWhile(lambda x: x < 10),
)
annotate(
"First ten after drop while <10",
nats.dropWhile(lambda x: x < 10).take(10),
)
public class Stream<A> {
private var _first : A
private var _rest : Stream<A>?
private var _step : () -> Stream<A>
init(_ first : A, _ step : @autoclosure () -> Stream<A>) {
self._first = first
self._rest = nil
self._step = step
}
public func first() -> A {
return _first
}
public func rest() -> Stream<A> {
if let stream = _rest {
return stream
} else {
let result = _step()
_rest = result
return result
}
}
public class func cons(head : A, _ tail: @autoclosure () -> Stream<A>) -> Stream<A> {
return Stream(head, tail)
}
public class func iterate(x : A, _ f : A -> A) -> Stream<A> {
return Stream(x, Stream.iterate(f(x), f))
}
public class func count(start : Int = 0, step : Int = 1) -> Stream<Int> {
return Stream<Int>.iterate(start) { $0 + step }
}
public class func repeat(x : A) -> Stream<A> {
return Stream(x, Stream.repeat(x))
}
public func map<B>(f : A -> B) -> Stream<B> {
return Stream<B>(f(self.first()), self.rest().map(f))
}
public func filter(predicate : A -> Bool) -> Stream<A> {
let x = self.first()
if predicate(x) {
return Stream(x, self.rest().filter(predicate))
} else {
return self.rest().filter(predicate)
}
}
public func zipWith<B, C>(other : Stream<B>, _ f : (A, B) -> C) -> Stream<C> {
return Stream<C>(
f(self.first(), other.first()),
self.rest().zipWith(other.rest(), f)
)
}
public func zip<B>(other : Stream<B>) -> Stream<(A, B)> {
return self.zipWith(other) { ($0, $1) }
}
public func scan<B>(first : B, _ f : (B, A) -> B) -> Stream<B> {
return Stream<B>(first, self.rest().scan(f(first, self.first()), f))
}
public func scan1(f : (A, A) -> A) -> Stream<A> {
return self.rest().scan(self.first(), f)
}
public func take(n : Int) -> [A] {
var i = 0
var result : [A] = []
var stream = self
while i < n {
result.append(stream.first())
stream = stream.rest()
i += 1
}
return result
}
public func drop(n : Int) -> Stream<A> {
var i = 0
var stream = self
while i < n {
stream = stream.rest()
i += 1
}
return stream
}
public func takeWhile(predicate : A -> Bool) -> [A] {
var result : [A] = []
var stream = self
while predicate(stream.first()) {
result.append(stream.first())
stream = stream.rest()
}
return result
}
public func dropWhile(predicate : A -> Bool) -> Stream<A> {
var stream = self
while predicate(stream.first()) {
stream = stream.rest()
}
return stream
}
}
public class RecursiveStream<A> : Stream<A> {
override init(_ first : A, _ step : @autoclosure () -> Stream<A>) {
super.init(first, step)
}
public class func forward(head : A) -> RecursiveStream<A> {
return RecursiveStream(head, Stream.repeat(head))
}
public func tie(step : @autoclosure () -> Stream<A>) {
self._step = step
}
}
func stars(n : Int) -> String {
let star : Character = "*"
return String(count: n, repeatedValue: star)
}
let nats = Stream.iterate(0) { $0 + 1 }
let ones = Stream.repeat(1)
let strings = nats.map { stars($0) }
let squares = nats.map { $0 * $0 }
let evens = nats.filter { $0 % 2 == 0 }
let odds = nats.filter { $0 % 2 == 1 }
let sums = nats.zipWith(nats) { $0 + $1 }
let pairs = evens.zip(odds)
/* Type inference for facts blows up if we don't specify <Int>
stream.swift:152:37: error: cannot invoke '*' with an argument of type
(($T8, ($T8, $T9) -> ($T8, $T9) -> $T7) -> ($T8, ($T8, $T9) -> $T7) ...
*/
let facts = Stream<Int>.count(start:1).scan1 { $0 * $1 }
/* Cannot use fibs inside of its own definition:
stream.swift:166:45: error: variable used within its own initial value
Fake with with
stream = RecursiveStream.forward()
stream.tie(...)
*/
let fibs = RecursiveStream.forward(0)
fibs.tie(Stream.cons(1, fibs.zipWith(fibs.drop(1)) { $0 + $1 }))
func annotate<A : Printable>(description : String, body : () -> A) {
println("\(description):\n \(body())\n")
}
annotate("First ten nats (iterate)") {
nats.take(10)
}
annotate("First ten nats after dropping ten nats (drop)") {
nats.drop(10).take(10)
}
annotate("First ten ones (repeat)") {
ones.take(10)
}
annotate("Type-changing map (map int -> string)") {
strings.take(10)
}
annotate("First ten squares (map int -> int)") {
squares.take(10)
}
annotate("First ten evens (filter int -> bool)") {
evens.take(10)
}
annotate("First ten odds (filter int -> bool)") {
odds.take(10)
}
annotate("First ten sums ((int, int) -> int)") {
sums.take(10)
}
annotate("Even/odd pairs (zip)") {
pairs.take(10)
}
annotate("First ten factorial (scan)") {
facts.take(10)
}
annotate("First ten fibs (forward/tie/zipwith)") {
fibs.take(10)
}
annotate("Take while <10") {
nats.takeWhile { $0 < 10 }
}
annotate("First ten after drop while <10") {
(nats.dropWhile { $0 < 10 }).take(10)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.