Address review comments better