Simplifying generator state, minor refactoring