Fix failing compilation

After reorganizing the code it was failing due to missing imports, this patch reintroduces them.