Simplify logic and fix bug (or two).

Thank your for the idea of this extension. I liked it so much I decided to hack on it!