A couple of years ago, I created a programming language. It was an attempt to make an extremely minimal logic programming system, with almost no built-in functionality other than the ability to do first-order logic, for the purposes of following the formalism in some more maths-driven textbooks. Recently, having learned Rust, I thought it would be a good exercise to port this language, originally written in Ruby, as I believe building it in Rust will result it significant performance wins and help make it more usable.

Though I’ve been making good progress with it, a few things have been surprisingly hard. One of the problems I’ve run into is that, while Rust’s iterators are highly expressive, performant and lazy, which helps substantially in this problem space, they run into problems when you try to model infinite sets with them. It turns out this is a significant problem in logic programming, and so I needed to find a good solution. So, this is a story about maths, infinities, and how computers deal with them, but first, we need to understand the problem logic programming languages try to solve.

Logical rules become tricky to deal with when they contain recursive references
to themselves. As an example, here’s a definition of the *subtyping relation*,
as given in Types and Programming Languages (TAPL). The expression `A <: B`

means that A is a subtype of B. One important property that this relation has is
that it’s *reflexive*, which means that any type is a subtype of itself for the
purposes of a type-checker. In the language I’m building, this rule is written
as follows:

```
rule S-Refl { $T <: $T }
```

`S-Refl`

is the name that TAPL gives this rule, and the repeated variable `$T`

means this rule matches statements where the expressions on either side of the
`<:`

symbol are the same.

Second, the relation is *transitive*, which means that if some type `S`

is a
subtype of `U`

, and that type in turn is a subtype of `T`

, then we can say that
`S`

is a subtype of `T`

. The rule below states that if `S <: U`

, and `U <: T`

,
then `S <: T`

, which is called `S-Trans`

in TAPL.

```
rule S-Trans {
$S <: $U / $U <: $T
-------------------
$S <: $T
}
```

These two rules don’t have any idea what `S`

, `T`

or `U`

are, they simply give
constraints that two types must obey in order to be considered subtypes. We’ll
now add some concrete facts to these rules, stating that gerbils are a subtype
of rodents, rodents a subtype of mammals, and mammals a subtype of animals.

```
rule S-Gerbil { Gerbil <: Rodent }
rule S-Rodent { Rodent <: Mammal }
rule S-Mammal { Mammal <: Animal }
```

Now, it ought to be possible to use these rules to generate a proof that gerbils are a subtype of animals. Indeed by combining the above definitions, we can construct such a proof:

```
──────────────── S-Rodent ──────────────── S-Mammal
Rodent <: Mammal Mammal <: Animal
──────────────── S-Gerbil ──────────────────────────────────────────── S-Trans
Gerbil <: Rodent Rodent <: Animal
────────────────────────────────────────────────────────── S-Trans
Gerbil <: Animal
```

The name to the right of each horizontal line indicates which rule produced that
part of the proof; the statements above the line are the rule’s premises (if
any) and the expression below the line is its conclusion. For example, the last
two lines of this proof can be read as “whereas `Gerbil <: Rodent`

, and ```
Rodent
<: Animal
```

, therefore by the rule `S-Trans`

we infer that `Gerbil <: Animal`

”.

A popular language for doing this kind of programming is called Prolog. In Prolog, the above logical rules can be directly transcribed as follows.

```
% S-Refl
subtype(T, T).
% S-Trans
subtype(S, T) :- subtype(S, U), subtype(U, T).
subtype(gerbil, rodent).
subtype(rodent, mammal).
subtype(mammal, animal).
```

Now, let’s say we want to find all the things of which `gerbil`

is a subtype;
we’re expecting the results `gerbil`

, `rodent`

, `mammal`

and `animal`

. To do
this computation, we give Prolog a query:

```
?- subtype(gerbil, X).
```

The Prolog interpreter will try to find all values of `X`

that make this a true
statement, that is, a statement matching the rules it’s been given. It starts
with the first rule: `subtype(T, T)`

. It compares the expressions
`subtype(gerbil, X)`

and `subtype(T, T)`

, and figures out what must be true in
order for them to be equal. For this to be the case, the sub-expressions
inside the parentheses must be equal, and this tells us that `T`

= `gerbil`

and `T`

= `X`

, therefore `X`

= `gerbil`

. And so it’s found our first solution to
`subtype(gerbil, X)`

.

Now, there are other results that can be found, by examining the other rules.
The next rule in the program (`S-Trans`

) begins with `subtype(S, T)`

, and
matching this against the query `subtype(gerbil, X)`

tells us `S`

= `gerbil`

and
`T`

= `X`

. Having matched the start of the rule, we try to match its first
premise, `subtype(S, U)`

. We know `S`

= `gerbil`

, so we now need to look for
rules that match `subtype(gerbil, U)`

. But this is just the query we started
with! So, from the first rule, we determine that `U`

= `gerbil`

. Similarly,
matching the second premise `subtype(U, T)`

, where we now know that `U`

=
`gerbil`

, gives us the same problem yet again and tells us that `T`

= `gerbil`

.
And so, this again tells us that `subtype(gerbil, gerbil)`

, with the following
proof:

```
──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
```

Can we generate any more results? In Prolog, the answer is unfortunately that we
cannot. Prolog’s search for results is *depth-first*; it will generate all the
possible results of looking at the `S-Trans`

rule, before moving on to any
others. So far it’s generated one result, by using `S-Refl`

to match both the
premises, and next it will try to generate more results for the first premise,
`subtype(S, U)`

, to see if they produce valid results. Well, having exhausted
the possibilities of `S-Refl`

, it moves onto `S-Trans`

, and so `S-Trans`

is now
recursively calling itself. This loop will continue indefinitely, so that Prolog
returns an infinite stream of `X = gerbil`

for our initial query. The next proof
it generates will be:

```
──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
────────────────────────────────────────── S-Trans ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
──────────────────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
```

After that it will just continue to infinitely recurse into the leftmost branch of this tree, never exploring other parts of the search space.

```
──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
────────────────────────────────────────── S-Trans ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
──────────────────────────────────────────────────────── S-Trans ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
──────────────────────────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
```

The core problem here is that there are various rules that could produce matches
for our query, but because one of them produces an infinite stream of results,
we never visit the later rules and generate more interesting results. In order
to reach those results, we can approach the problem *breadth-first* by
generating one result for each rule, before looking for further results that any
of the rules produce. This is the approach taken by the μKanren language.
In this system, the first three proofs of `subtype(gerbil, X)`

would be:

```
1. ──────────────── S-Refl
Gerbil <: Gerbil
2. ──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
3. ──────────────── S-Gerbil
Gerbil <: Rodent
```

We’ve already generated a result we couldn’t reach before: since we didn’t get
stuck looking at results from `S-Trans`

, we can visit `S-Gerbil`

and later rules
to see which ones match our query. Only after this first pass do we return to
any of the rules and see if they produce more matches. `S-Refl`

and `S-Gerbil`

produce one result each, since they have no premises, but `S-Trans`

will produce
more. We can imagine using the above three results in place of this rule’s first
premise, `subtype(S, U)`

where `S`

= `gerbil`

. Using result 1 actually generates
result 2 above, but using result 2 gives us the next proof:

```
4. ──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
────────────────────────────────────────── S-Trans ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
──────────────────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
```

For each choice of the first premise, we also want to visit different choices
for the second premise, and the next result we generate combines `S-Refl`

for
the first premise with `S-Trans`

for the second:

```
5. ──────────────── S-Refl ──────────────── S-Refl
Gerbil <: Gerbil Gerbil <: Gerbil
──────────────── S-Refl ────────────────────────────────────────── S-Trans
Gerbil <: Gerbil Gerbil <: Gerbil
─────────────────────────────────────────────────────── S-Trans
Gerbil <: Gerbil
```

We can also use `S-Gerbil`

for the first premise and `S-Relf`

for the second:

```
6. ──────────────── S-Gerbil ──────────────── S-Refl
Gerbil <: Rodent Rodent <: Rodent
──────────────────────────────────────────── S-Trans
Gerbil <: Rodent
```

These proofs aren’t generating any new *conclusions* yet, but they will
eventually, as we’re exploring more branches of the search space than Prolog
does when it recurses off down the left branch indefinitely. After generating
more combinations of the above proofs, we’ll eventually generate this, where
`S-Trans`

uses two of the later rules for its premises:

```
──────────────── S-Gerbil ──────────────── S-Rodent
Gerbil <: Rodent Rodent <: Mammal
──────────────────────────────────────────── S-Trans
Gerbil <: Mammal
```

Given enough time, the whole search space will be explored, and it’s simply a matter of waiting for a result you’re interested in, rather than never being able to generate it.

This problem of enumerating infinite sets is illustrated by the thought experiment known as Hilbert’s hotel. It describes an imaginary hotel that has an infinite number of rooms. One day, an infinite number of buses, each carrying an infinite number of passengers arrive at the hotel, and all the passengers need a place to stay the night. The problem is how to make sure all the visitors are given a room.

Surely, after the first infinite busload of visitors disembark and are allocated
rooms, the hotel is now full, and all the other buses must be turned away. This
is equivalent to following only the infinite stream of results from `S-Trans`

and never visiting any of the other rules. Let’s illustrate this situation by
representing the rooms as an array of boxes, and the number in each box
identifies which bus, and which seat on that bus, the guest came from.

```
+-----+-----+-----+-----+-----+ ...
| 1.1 | 1.2 | 1.3 | 1.4 | 1.5 |
+-----+-----+-----+-----+-----+ ...
```

However, it turns out we can accommodate at least one more busload of guests. Once the hotel is “full” of the first bus’s passengers, we ask all those guests to move into the room whose number is double the number of their current room. The guest in room 1 moves to room 2, the one in room 2 to room 4, and so on. The hotel being infinite, it’s possible to move all the guests in this way and still leave all the odd-numbered rooms empty.

```
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
| | 1.1 | | 1.2 | | 1.3 | | 1.4 | | 1.5 |
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
```

The next busload of visitors can then move into the odd-numbered rooms, of which there is an infinite quantity, and the hotel is again full.

```
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
| 2.1 | 1.1 | 2.2 | 1.2 | 2.3 | 1.3 | 2.4 | 1.4 | 2.5 | 1.5 |
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
```

What about all the other infinite buses waiting on the infinite driveway to unload their passengers? Well for each bus that arrives we can repeat this process, moving all the existing guests into the even-numbered rooms and letting the next bus fill up the odd numbers. And so, the infinite queue of infinitely large buses can all be accommodated.

```
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
| | 2.1 | | 1.1 | | 2.2 | | 1.2 | | 2.3 |
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
| 3.1 | 2.1 | 3.2 | 1.1 | 3.3 | 2.2 | 3.4 | 1.2 | 3.5 | 2.3 |
+-----+-----+-----+-----+-----+-----+-----+-----+-----+-----+ ...
```

This is a neat illustration of the paradoxes of infinity: that there are the same number of even (or odd) numbers as there are integers, and that an infinite set of infinite sets (the buses and their passengers) is no larger than a single infinite set (the hotel’s rooms). But it also gives us a clue to solve our infinite search problem.

A logic interpreter works by trying to match a given query against the
conclusions of all the known rules. Matching a pair of expressions generates a
*state*, a set of bindings recording the variables that were found in the
expressions and what they’ve been bound to. For example, matching
`subtype(gerbil, X)`

against `subtype(T, T)`

, with no prior state, generates the
state `{ T = gerbil, X = gerbil }`

. For each matching rule, we then try to match
the rule’s first premise, in the state generated by the conclusion. This match
may generate an infinite number of results, and thus an infinite stream of
matching states. In our example above, we saw that the first premise of
`S-Trans`

, `subtype(S, U)`

, generated such a stream.

The search proceeds by matching the next premise in each of these states, in
principle generating an infinite stream of infinite streams of states. For each
state generated by the premise `subtype(S, U)`

, the next premise `subtype(U, T)`

may also generate an infinite stream of states. And so the following premise, if
there is one, needs to be checked in any of an infinite set of infinite sets of
states, exactly the problem illustrated by Hilbert’s hotel.

Prolog gets stuck because it only looks at the set generated by the first premise, which never ends, and so it never moves on to consider the state space generated by the other premises, or other rules. It only lets the passengers from the first bus into the hotel. It works depth-first, whereas we need a breadth-first search to have any hope of reaching other states. However, it’s not as simple as that; a naive breadth-first search of a stream of streams would mean taking the first item from each stream, then the second, and so on. But if the outer stream is infinite, we can never finish taking the first elements of the inner streams, and the second elements are never visited. The first passenger on each bus gets a room, and everyone else must sleep in their seats.

To solve the problem, we need to *interleave* the sets. That’s what we were
doing when we moved the existing guests into the even-numbered rooms in order to
let the next busload into the hotel. Now, before looking at how to do this in
Rust, it’s instructive to look at how it’s done in Haskell, which is probably
the easiest language in which to implement this. Whereas Rust has lazy
iterators, Haskell has lazy *everything*; all parts of Haskell programs are only
evaluated when strictly necessary to produce some output. And like Rust, Haskell
has infinite iterators: the notation `[1..]`

means an infinite list of all the
integers, starting from 1. We can’t print this whole list, but just like in Rust
we can use `take`

to get only the first N elements of it – `take 5 [1..]`

gives
the finite list `[1, 2, 3, 4, 5]`

.

Using `map`

, we can generate an infinite list of infinite lists. ```
map (\n ->
[(100 * n)..]) [0..]
```

is equivalent to the Rust expression ```
(0..).map(|n| (100 *
n)..)
```

, and produces an infinite list of the ranges `[0..]`

, `[100..]`

,
`[200..]`

, and so on.

First we’ll look at two ways of accommodating the guests that don’t work. One way is to concatenate all the buses’ sets of passengers, but, as the first bus is infinite, we never finish admitting its passengers and all the other buses remain blocked. Here that’s demonstrated by only getting elements from the first range in the output, equivalent to a depth-first search.

```
let buses = map (\n -> [(100 * n)..]) [0..] in
print (take 32 (concat buses))
-- -> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
-- 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
```

The other simple approach is to use a breadth-first search: take the first
passenger from each bus, then the second, and so on. However, as there are an
infinite number of buses, we never complete the first iteration, and we only get
the first member of each set. In Haskell, this is equivalent to ```
concat
(transpose buses)
```

.

```
let buses = map (\n -> [(100 * n)..]) [0..] in
print (take 32 (concat (transpose buses)))
-- -> [0, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200,
-- 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300,
-- 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100]
```

To properly solve the problem, we need a function for interleaving two lists. In
Haskell, lists can be destructured into their first element (`head`

) and the
rest of the list (`tail`

) using the `:`

operator, so `(x:xs)`

pattern-matches a
list of at least one element, binding that element to `x`

and the rest of the
list to `xs`

. `interleave`

takes two lists of any type (`[a]`

) and returns
another list. If the first list is empty, then the result of interleaving it
with the second is just the second list. If the first list has at least one
element, then the result is that element, followed by the result of interleaving
the second list with the rest of the first – swapping the arguments is what
achieves the interleaving effect.

```
interleave :: [a] -> [a] -> [a]
interleave [] ys = ys
interleave (x:xs) ys = x : interleave ys xs
```

To flatten an infinite list of infinite lists (`[[a]]`

) into a single list
(`[a]`

), we interleave the first inner list with the result of flattening the
rest of the lists.

```
flatten :: [[a]] -> [a]
flatten [] = []
flatten (x:xs) = interleave x (flatten xs)
```

This is a tiny amount of code, but Haskell’s lazy evaluation means we can directly write the mathematical definition of what we’re doing, and it will compile and run and give a meaningful answer:

```
let buses = map (\n -> [(100 * n)..]) [0..] in
print (take 32 (flatten buses))
-- -> [0, 100, 1, 200, 2, 101, 3, 300, 4, 102, 5, 201, 6, 103, 7, 400, 8,
-- 104, 9, 202, 10, 105, 11, 301, 12, 106, 13, 203, 14, 107, 15, 500]
```

The numbers in the result are not evenly distributed; every other element is taken from the first set. But that’s not so important as the fact that we are visiting every input list, rather only the first list, or only getting the first member of each list. Given an infinite amount of time, all the input lists will be expanded – all the buses can unload all their passengers.

In Haskell, this is possible because of how laziness relies on recursive
definitions. `interleave`

is defined in terms of itself, and `flatten`

is
defined in terms of itself and `interleave`

. And the particular recursive
structure of the definitions is what makes them work. To make lazy lists work,
all we need to be able to do is find the next element, without evaluating the
entire list. So say we need to flatten this list:

```
buses = [[0..], [100..], [200..], ..]
```

The flattened version of `buses`

is given by this expression:

```
flatten buses
```

Can we determine just the first element of this, without evaluating the entire
result? Well, let’s look at the definition of `flatten`

, which tells us that
`flatten (x:xs) = interleave x (flatten xs)`

. Dropping `buses`

into this
equation gives:

```
flatten buses = interleave [0..] (flatten [[100..], [200..], ..])
```

Now we look at the definition of `interleave`

, which says ```
interleave (x:xs) ys
= x : interleave ys xs
```

. The infinite list `[0..]`

destructures into ```
0 :
[1..]
```

, and so applying this to the above gives:

```
flatten buses = 0 : interleave (flatten [[100..], [200..], ..]) [1..]
```

We now know the first element of the result: `0`

! Let’s try one more round of
substitutions:

```
0 : interleave (flatten [[100..], [200..], ..]) [1..]
-- destructure [[100..], [200..], ..] into [100..] : [[200..], ..]
= 0 : interleave (flatten ([100..] : [[200..], ..])) [1..]
-- apply definition flatten (x:xs) = interleave x (flatten xs)
= 0 : interleave (interleave [100..] (flatten [[200..], ..])) [1..]
-- destructure [100..] into 100 : [101..]
= 0 : interleave (interleave (100 : [101..]) (flatten [[200..], ..])) [1..]
-- apply definition interleave (x:xs) ys = x : interleave ys xs (inner)
= 0 : interleave (100 : interleave (flatten [[200..], ..]) [101..]) [1..]
-- apply definition interleave (x:xs) ys = x : interleave ys xs (outer)
= 0 : 100 : interleave [1..] (interleave (flatten [[200..], ..]) [101..])
```

So now we know that the result begins `[0, 100, ..]`

, without having to evaluate
any further. The `take`

function essentially does this until it’s generated the
requested number of results. We can see elements being popped off the front of
the inner ranges, and “pulled” through the definitions of `flatten`

and
`interleave`

as they pattern-match the first element of their inputs to produce
a new expression.

If you’re familiar with Rust, this should remind you of something: the
`Iterator`

trait. This trait defines the `next()`

method, which pulls
values out of the iterator on demand, allowing iterators to be composed lazily.
However, while `Iterator`

itself has a lazy interface, Rust in general is
eagerly evaluated, and that’s going to pose a challenge for our implementation.
We’ll also incur a fair amount of boilerplate as we convert between various
types, so we’ll need a lot more code than we did in Haskell. But, once we have
the right iterators defined, we can compose them in powerful ways.

Looking at the Haskell evaluation above, it appears there are three fundamental iterators present that we need to turn into Rust equivalents.

Sources, i.e. the actual streams of values that we want to combine; in our Haskell values these are infinite lists of ints, but they could be anything that implements

`Iterator<Item = T>`

.Interleaves, which take two (or more) streams and return the next element by alternating between the inputs.

Flattens, which take a stream of streams, and return the next element by interleaving the first element with the result of flattening the rest.

We’ll start by defining a type alias. Since `Iterator`

is a trait and we’ll need
to store values that implement it, we need to create trait objects by boxing
things. A `BoxIter`

is a boxed iterator of some element type `T`

.

```
type BoxIter<'a, T> = Box<dyn Iterator<Item = T> + 'a>;
```

First, let’s define `Interleave`

. Although in Haskell this was defined
recursively, in Rust it’s possible to define using a loop; the struct contains a
fixed and finite number of iterators, and a `cursor`

that points to the next one
it’s going to try to pull from.

```
struct Interleave<'a, T> {
cursor: usize,
iters: Vec<BoxIter<'a, T>>,
}
```

We can create one of these structs given a set of iterators. To accept as broad
a range of inputs as possible, we’ll make `Interleave::new`

accept any
`IntoIterator`

value, whose items are themselves `IntoIterator`

values that
yield `T`

. The `new()`

function turns the input into an iterator, maps its
contents to boxed iterators, and collects the result into a `Vec`

, which is
stored in the returned `Interleave`

struct.

```
impl<'a, T> Interleave<'a, T> {
pub fn new<I, J>(iters: I) -> Interleave<'a, T>
where
I: IntoIterator<Item = J>,
J: IntoIterator<Item = T> + 'a,
{
let iters: Vec<_> = iters
.into_iter()
.map(|s| Box::new(s.into_iter()) as BoxIter<'a, T>)
.collect();
let cursor = if iters.is_empty() { 0 } else { iters.len() - 1 };
Interleave { cursor, iters }
}
}
```

The `Iterator`

implementation for this struct is relatively straightforward; its
role is to pull items from each of the iterators in turn. So, starting from the
current `cursor`

index, we call `next()`

on each inner iterator until one of
them returns `Some`

. If we try every iterator and none of them returns anything,
then we return `None`

.

```
impl<T> Iterator for Interleave<'_, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
for _ in 0..self.iters.len() {
self.cursor = (self.cursor + 1) % self.iters.len();
let item = self.iters[self.cursor].next();
if item.is_some() {
return item;
}
}
None
}
}
```

This is enough to produce breadth-first behaviour, i.e. picking the first
passenger from each bus. Note that we need to restrict `buses`

to a finite set
here using `take()`

, so that we can create an `Interleave`

from them.

```
let buses = (0..).map(|n| (100 * n)..).take(1_000);
let interleaved: Vec<_> = Interleave::new(buses).take(32).collect();
println!("{:?}", interleaved);
// -> [0, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200,
// 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300,
// 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100]
```

In contrast, Rust’s built-in `flatten()`

method produces depth-first behaviour,
taking elements only from the first stream and not from any of the others. Here
`buses`

can be an infinite iterator, and iterating it for a while shows that it
only takes elements from the first range.

```
let buses = (0..).map(|n| (100 * n)..);
let flat: Vec<_> = buses.flatten().take(32).collect();
println!("{:?}", flat);
// -> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
// 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
```

Now we can tackle `Flatten`

. Let’s remind ourselves of the Haskell definition of
this iterator:

```
flatten :: [[a]] -> [a]
flatten [] = []
flatten (x:xs) = interleave x (flatten xs)
```

Given some iterator `(x:xs)`

, it takes the first element from it and interleaves
that with the result of flattening the rest of the iterator. And so, it might be
tempting to try to implement it like this:

```
fn flatten<'a, I, J, T>(mut streams: I) -> Option<Interleave<'a, T>>
where
I: Iterator<Item = J> + 'a,
J: Iterator<Item = T> + 'a,
T: 'a,
{
if let (Some(head), Some(tail)) = (streams.next(), flatten(streams)) {
Some(Interleave::new(vec![
Box::new(head) as BoxIter<'a, T>,
Box::new(tail) as BoxIter<'a, T>,
]))
} else {
None
}
}
```

However, the unconditional recursive call to `flatten()`

means that if `streams`

is infinite, this function will cause a stack overflow. Rust iterators are lazy
in that they call `next()`

only when more elements are requested, but the code
for *constructing* them is not lazy, it’s just ordinary Rust code. To get
laziness, we need to defer the creation of the `tail`

iterator here until we
actually need to get the next item from it.

So, we need to create a distinct iterator type that uses a closure to defer the
creation of that inner `Interleave`

struct. That closure will be called as
needed to generate a value, which we’ll denote by wrapping it in a `Thunk`

,
another type we’ll define shortly. The closure will return an
`Option<Interleave<T>>`

; it’s optional because we’ll only generate one if the
stream we’re given has a first element.

```
struct Flatten<'a, T> {
thunk: Thunk<'a, Option<Interleave<'a, T>>>,
}
impl<'a, T> Flatten<'_, T> {
fn new(mut streams: BoxIter<'a, BoxIter<'a, T>>) -> Flatten<'a, T> {
let thunk = Thunk::new(|| {
streams.next().map(|head| {
let tail = Box::new(Flatten::new(streams));
Interleave::new(vec![head, tail])
})
});
Flatten { thunk }
}
}
```

Implementing `Iterator`

for this struct requires getting a mutable reference to
the value inside the `Thunk`

, which we can do using the `AsMut`

trait; this
will return `&mut Option<Interleave<T>>`

, and we can pattern-match on that to
return the next item from the wrapped stream, if it exists.

```
impl<T> Iterator for Flatten<'_, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if let Some(stream) = self.thunk.as_mut() {
stream.next()
} else {
None
}
}
}
```

The trick that makes this all work is the `Thunk`

type. Its job is to defer the
creation of a value by wrapping a closure that generates it. We only want to
generate the inner value once, and so the `Thunk`

exists in one of two states:
`Pending`

, which contains a boxed `FnOnce`

that returns a `T`

, or `Forced`

,
which contains a `T`

. The `Option`

inside `Pending`

is a trick that lets us move
the closure out of the struct, which is required if we want to call it. Because
it’s wrapping a non-nullable pointer, the `Option`

does not require any space to
be allocated, it’s just there to make the borrow checker happy.

```
enum Thunk<'a, T: 'a> {
Pending(Option<Box<dyn 'a + FnOnce() -> T>>),
Forced(T),
}
impl<'a, T> Thunk<'a, T> {
fn new<F>(gen: F) -> Thunk<'a, T>
where
F: 'a + FnOnce() -> T,
{
Thunk::Pending(Some(Box::new(gen)))
}
}
```

Implementing `AsMut`

for `Thunk`

requires us to return `&mut T`

. To do this, we
need to evaluate the closure if it’s not already been used; if the thunk is
`Pending`

, then we can `take()`

the closure out and call it, generating a
`T`

, which we store as a `Forced`

value. Once this is done, we can assume the
thunk is in the `Forced`

state, so we pattern-match the value out of it and
return a mutable reference to it.

```
impl<T> AsMut<T> for Thunk<'_, T> {
fn as_mut(&mut self) -> &mut T {
if let Thunk::Pending(gen) = self {
let gen = gen.take().unwrap();
*self = Thunk::Forced(gen());
}
match self {
Thunk::Forced(ref mut value) => value,
_ => panic!(),
}
}
}
```

This interface lets us get a mutable reference to the iterator inside the thunk,
and thus ask it for the next item. Finally, we can add a little syntactic sugar
so that we can build a `Flatten`

from any nested `IntoIterator`

value – just as
we did for `Interleave`

, this requires converting the inputs into boxed
iterators so we can pass them to the normal `Flatten`

constructor.

```
impl<'a, T> Flatten<'a, T> {
pub fn from_iter<I, J>(streams: I) -> Flatten<'a, T>
where
I: IntoIterator<Item = J> + 'a,
J: IntoIterator<Item = T> + 'a,
{
let iters = streams
.into_iter()
.map(|s| Box::new(s.into_iter()) as BoxIter<'a, T>);
Flatten::new(Box::new(iters))
}
}
```

With these three structs defined – `Interleave`

, `Thunk`

and `Flatten`

– we
can convert our infinite list of infinite lists into a flattened form that pulls
items from every iterator in the set:

```
let buses = (0..).map(|n| (100 * n)..);
let flat: Vec<_> = Flatten::from_iter(buses).take(32).collect();
println!("{:?}", flat);
// -> [0, 100, 1, 200, 2, 101, 3, 300, 4, 102, 5, 201, 6, 103, 7, 400, 8,
// 104, 9, 202, 10, 105, 11, 301, 12, 106, 13, 203, 14, 107, 15, 500]
```

And so we’ve reproduced the behaviour that took a handful of lines in Haskell.
Now, maybe I’m using the wrong language for this project, but there’s also been
a good deal of it that Rust’s lazy iterators made substantially easier. The
thing to remember is that if something returns an iterator – think `map`

,
`filter`

, or `take`

– then it may be lazy. But anything that
returns a non-iterator value must be eager. The place I tripped over this the
most was when using `fold`

, especially since I was using it to transform
one iterator into another. I expected the closure to be called lazily, but it
wasn’t, and that forced the creation of all of the above. But once you have the
building blocks you need, Rust’s iterators are remarkably useful and compose
wonderfully, giving you expressive code without the overhead of many nested
loops as you’d get if they were eager.