Cleaned up code to use state monad where appropraite rather than faking it by threading a mutable reference.