Fix handling of state. Optimize a bit.