Fixing code structure