Code cleanup after merge