Fix bugs and improve error message from previous commit