Code cleanup (amended)