Support for in-place return in eval fn. (#6954)

Ignore storage arguments when evaluating a call, like we do for other
kinds of instruction. Create a placeholder constant to represent each
out parameter so that it can be used in the function body to form more
storage arguments.

Assisted-by: Gemini 3.1 Pro via Antigravity

toolchain/check/eval.cpp

@@ -588,24 +588,29 @@ static auto GetConstantValue(EvalContext& eval_context,
                              Phase* phase) -> SemIR::InstBlockId = delete;
 
 // If the given instruction block contains only constants, returns a
-// corresponding block of those values.
-static auto GetConstantValue(EvalContext& eval_context,
-                             SemIR::InstBlockId inst_block_id, Phase* phase)
-    -> SemIR::InstBlockId {
+// corresponding block of those values. Ignores the instructions in the
+// specified range of indexes, replacing those elements with `None`.
+static auto GetConstantBlockValueIgnoringIndexRange(
+    EvalContext& eval_context, SemIR::InstBlockId inst_block_id, Phase* phase,
+    std::pair<int, int> ignored_range) -> SemIR::InstBlockId {
   if (!inst_block_id.has_value()) {
     return SemIR::InstBlockId::None;
   }
   auto insts = eval_context.inst_blocks().Get(inst_block_id);
   llvm::SmallVector<SemIR::InstId> const_insts;
   for (auto inst_id : insts) {
-    auto const_inst_id = GetConstantValue(eval_context, inst_id, phase);
-    if (!const_inst_id.has_value()) {
-      return SemIR::InstBlockId::None;
+    auto const_inst_id = SemIR::InstId::None;
+    if (static_cast<int>(const_insts.size()) < ignored_range.first ||
+        static_cast<int>(const_insts.size()) >= ignored_range.second) {
+      const_inst_id = GetConstantValue(eval_context, inst_id, phase);
+      if (!const_inst_id.has_value()) {
+        return SemIR::InstBlockId::None;
+      }
     }
 
     // Once we leave the small buffer, we know the first few elements are all
-    // constant, so it's likely that the entire block is constant. Resize to the
-    // target size given that we're going to allocate memory now anyway.
+    // constant, so it's likely that the entire block is constant. Resize to
+    // the target size given that we're going to allocate memory now anyway.
     if (const_insts.size() == const_insts.capacity()) {
       const_insts.reserve(insts.size());
     }
@@ -617,6 +622,15 @@ static auto GetConstantValue(EvalContext& eval_context,
   return eval_context.inst_blocks().AddCanonical(const_insts);
 }
 
+// If the given instruction block contains only constants, returns a
+// corresponding block of those values.
+static auto GetConstantValue(EvalContext& eval_context,
+                             SemIR::InstBlockId inst_block_id, Phase* phase)
+    -> SemIR::InstBlockId {
+  return GetConstantBlockValueIgnoringIndexRange(eval_context, inst_block_id,
+                                                 phase, {0, 0});
+}
+
 // Compute the constant value of a type block. This may be different from the
 // input type block if we have known generic arguments.
 static auto GetConstantValue(EvalContext& eval_context,
@@ -2088,6 +2102,39 @@ static auto TryEvalCall(EvalContext& outer_eval_context, SemIR::LocId loc_id,
                         SemIR::SpecificId specific_id,
                         SemIR::InstBlockId args_id) -> SemIR::ConstantId;
 
+// Returns the range of parameter indexes that contain the return storage for
+// this function call.
+static auto GetReturnStorageParamIndexRange(EvalContext& eval_context,
+                                            const SemIR::Callee& callee)
+    -> std::pair<int, int> {
+  if (const auto* callee_function =
+          std::get_if<SemIR::CalleeFunction>(&callee)) {
+    const auto& function =
+        eval_context.functions().Get(callee_function->function_id);
+    return {function.call_param_ranges.return_begin().index,
+            function.call_param_ranges.return_end().index};
+  }
+
+  return {0, 0};
+}
+
+// Replace the `args_id` field of a call with its constant value. The return
+// storage argument, if any, is instead replaced with `None`.
+static auto ReplaceCallArgsFieldWithConstantValue(EvalContext& eval_context,
+                                                  const SemIR::Callee& callee,
+                                                  SemIR::Call* call,
+                                                  Phase* phase) -> bool {
+  auto return_storage_param_index_range =
+      GetReturnStorageParamIndexRange(eval_context, callee);
+  auto args_id = GetConstantBlockValueIgnoringIndexRange(
+      eval_context, call->args_id, phase, return_storage_param_index_range);
+  if (!args_id.has_value() && call->args_id.has_value()) {
+    return false;
+  }
+  call->args_id = args_id;
+  return IsConstantOrError(*phase);
+}
+
 // Makes a constant for a call instruction.
 static auto MakeConstantForCall(EvalContext& eval_context,
                                 SemIR::InstId inst_id, SemIR::Call call)
@@ -2129,8 +2176,8 @@ static auto MakeConstantForCall(EvalContext& eval_context,
   bool has_constant_operands =
       has_constant_callee &&
       ReplaceTypeWithConstantValue(eval_context, inst_id, &call, &phase) &&
-      ReplaceFieldWithConstantValue(eval_context, &call, &SemIR::Call::args_id,
-                                    &phase);
+      ReplaceCallArgsFieldWithConstantValue(eval_context, callee, &call,
+                                            &phase);
   if (phase == Phase::UnknownDueToError) {
     return SemIR::ErrorInst::ConstantId;
   }
@@ -2837,6 +2884,24 @@ auto TryExecTypedInst<SemIR::OutParam>(FunctionExecContext& eval_context,
     return SemIR::ConstantId::None;
   }
 
+  if (!eval_context.args()[param.index.index].has_value()) {
+    // The argument will be `None` for an index corresponding to a return
+    // storage argument for return values that have an in-place initializing
+    // representation. Produce an opaque "out parameter" variable for now, so
+    // that references to it can still successfully evaluate.
+    //
+    // TODO: Create and track mutable storage for the return value here. This is
+    // necessary to support things like `returned var`.
+    eval_context.locals().Insert(
+        inst_id,
+        MakeConstantResult(
+            eval_context.context(),
+            SemIR::VarStorage{.type_id = inst.type_id(),
+                              .pattern_id = SemIR::AbsoluteInstId::None},
+            Phase::Concrete));
+    return SemIR::ConstantId::None;
+  }
+
   return TryExecTypedParam(eval_context, inst_id, inst);
 }