Fix a bunch of state restoration bugs