UPDATE: cleanup code