Code design fix