Updated the design of code