   return {s_result_buffer_.data(), s_result_buffer_.data() + n_available};
 }
 
-__global__ void QueryFeatureListKernel(common::Span<int32_t> feature_list_input,
-                                       common::Span<int32_t> node_feature_list,
-                                       BitField result_buffer_input,
-                                       BitField result_buffer_output) {
+__global__ void SetInputBufferKernel(common::Span<int32_t> feature_list_input,
+                                     BitField result_buffer_input) {
   uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
   if (tid < feature_list_input.size()) {
     result_buffer_input.Set(feature_list_input[tid]);
   }
-
-  if (tid < node_feature_list.size()) {
-    result_buffer_output.Set(node_feature_list[tid]);
-  }
+}
+
+__global__ void QueryFeatureListKernel(BitField node_constraints,
+                                       BitField result_buffer_input,
+                                       BitField result_buffer_output) {
+  result_buffer_output |= node_constraints;
   result_buffer_output &= result_buffer_input;
 }
 

}

__global__ void ClearBuffersKernel(
    BitField result_buffer_self, BitField result_buffer_input, BitField feature_buffer) {
  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid < result_buffer_self.Size()) {
    result_buffer_self.Clear(tid);
  }
  if (tid < result_buffer_input.Size()) {
    result_buffer_input.Clear(tid);

  if (!has_constraint_ || nid == 0) {
    return feature_list;
  }

  ClearBuffers();

  BitField node_constraints = s_node_constraints_[nid];
  CHECK_EQ(input_buffer_bits_.Size(), output_buffer_bits_.Size());

  int constexpr kBlockThreads = 256;
  const int n_grids = static_cast<int>(
      dh::DivRoundUp(output_buffer_bits_.Size(), kBlockThreads));
  SetInputBufferKernel<<<n_grids, kBlockThreads>>>(feature_list, input_buffer_bits_);

  QueryFeatureListKernel<<<n_grids, kBlockThreads>>>(
      node_constraints, input_buffer_bits_, output_buffer_bits_);

  thrust::counting_iterator<int32_t> begin(0);
  thrust::counting_iterator<int32_t> end(result_buffer_.size());
