Sorry for my previous, hasty patch.
On closer inspection, it's not hard to incorporate my changes with yours.
Please try out this new patch rather than the old one. It should work much better.

--Kyle