Some design fixes for issue 3