Fix missing imports

(oops, half-committed some work in progress changes)