Fixup and update implementation for the new design.
This will create a new commit in order to revert the existing changes.
