Address review comments