Rollback fix and code cleanup