Light code refactor