Generator rejig.
updated TODO