const int state_size =
      reduce_dataset->getAttrOfType<ArrayAttr>("Tstate").size();
  for (int i = 0; i < state_size; ++i) {
    body_returns.push_back(dataset_if.getResult(i));
  }
  // Copy the arguments but skip the states and the loop condition
  // which are updated in while body.
  for (int i = state_size + 1; i < body_args.size(); ++i) {