This would add some complexity to `set_state`, but it probably would be a lot nicer from the user perspective. I'll work on moving it that way.