Cleanup to stick to the architecture