Fun with Generators

Tagged:  •    •  

Sometimes I just wander around in blogs looking for cool snippets in other programming languages. Since the upcoming GNU Smalltalk 3.0 will have some features from Python and Ruby (most notably generators), it is nice to test corner cases and see if GNU Smalltalk's behavior matches what happens in other languages.

I came across two implementations of the power-set function in Python, a recursive and an iterative one:

# by Guido van Rossum
def power(s):
    if len(s) == 0:
        yield Set()
        # Non-destructively choose a random element:
        x = Set([iter(s).next()])
        for ss in power(s - x):
            yield ss
            yield ss | x

# by Tim Peters
def power(base):
    pairs = [(2**i, x) for i, x in enumerate(s)]
    for i in xrange(2**len(pairs)):
        yield [x for (mask, x) in pairs if i & mask]

Both change in interesting way when you translate them in Smalltalk. The difference is mostly because Python is being used as a procedural language here, while in Smalltalk you have to think about where to place your methods. Sometimes the resulting code seems less compact, but only because you add utility methods that could be useful elsewhere—so that ultimately the solution is more complete.

For example, in translating the first example I decide to implement #powerSet as a method on Stream (of which Generator is a subclass) and not on Collection. Using Streams has the advantage that we can use #next as a destructive way to do, in a single step, iter(s).next() and s - x. Then, we need a couple of tiny forwarding methods in Collection. One just forwards #power to a Stream, and the second adds a #readStream method to Collection, that uses generators to simply wrap #do: iteration. Like this:

Collection extend [
    readStream [
        ^(Generator on: [ :g | self do: [ :each | g yield: each ] ])

    power [
        ^self asSet readStream power

(side note: the first of these two methods might find its way into GNU Smalltalk's class library...)

Then the body of the main recursive method looks very similar to Guido's:

Stream extend [
    power [
        ^Generator on: [ :g |
            | some |
            self atEnd
                ifTrue: [ g yield: Set new ]
                ifFalse: [
                    some := self next.
                    self power do: [ :each |
                        g yield: each.
                        g yield: each + {o} ] ] ]

Of course, here is how you use them

#(3 4 5) power do: [ :each |
    each asArray printNl ]

Now, on to the iterative version. The code is very simple, possibly easier than in Python because we can use the little-known #bitAt: method:

Collection extend [
    power [
        ^(0 to: (1 bitShift: self size) - 1) readStream collect: [ :each || i |
            i := 0.
            self select: [ :elem | (each bitAt: (i := i + 1)) = 1 ] ]

Also note that I didn't use generators here. Instead, I rely on Stream #collect: to lazily evaluate the block and generate combination only as they are requested. That's it!

User login