Cleanup code; Add missing parameter