[ VM ] Ensure TypeArguments register is preserved when regenerating allocation stubs for parameterized classes

Fixes https://github.com/flutter/flutter/issues/88104

TEST=pkg/vm_service/test/regress_88104_test.dart

Change-Id: I87affc62189bc076cf6e46c47e76f4fb005f9068
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/243850
Reviewed-by: Alexander Markov <alexmarkov@google.com>
Reviewed-by: Ryan Macnak <rmacnak@google.com>

pkg/vm_service/test/regress_88104_test.dart

@@ -0,0 +1,48 @@
+// Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+// Regression test for https://github.com/flutter/flutter/issues/88104.
+//
+// Ensures that the `TypeArguments` register is correctly preserved when
+// regenerating the allocation stub for generic classes after enabling
+// allocation tracing.
+import 'dart:async';
+import 'dart:developer';
+
+import 'package:vm_service/vm_service.dart';
+
+import 'common/service_test_common.dart';
+import 'common/test_helper.dart';
+
+class Foo<T> {}
+
+testMain() async {
+  debugger();
+  for (int i = 0; i < 10; ++i) {
+    Foo<int>();
+    await Future.delayed(const Duration(milliseconds: 10));
+  }
+}
+
+final tests = <IsolateTest>[
+  hasStoppedAtBreakpoint,
+  (VmService service, IsolateRef isolateRef) async {
+    final isolateId = isolateRef.id!;
+    final isolate = await service.getIsolate(isolateId);
+    final rootLibId = isolate.rootLib!.id!;
+    final rootLib = await service.getObject(isolateId, rootLibId) as Library;
+    final fooCls = rootLib.classes!.first;
+    await service.setTraceClassAllocation(isolateId, fooCls.id!, true);
+  },
+  resumeIsolate,
+  hasStoppedAtExit,
+];
+
+main([args = const <String>[]]) => runIsolateTests(
+      args,
+      tests,
+      'regress_88104_test.dart',
+      testeeConcurrent: testMain,
+      pause_on_exit: true,
+    );

runtime/vm/compiler/stub_code_compiler_arm.cc

@@ -737,6 +737,33 @@ void StubCodeCompiler::GenerateFixAllocationStubTargetStub(
   __ Branch(FieldAddress(R0, target::Code::entry_point_offset()));
 }
 
+// Called from object allocate instruction when the allocation stub for a
+// generic class has been disabled.
+void StubCodeCompiler::GenerateFixParameterizedAllocationStubTargetStub(
+    Assembler* assembler) {
+  // Load code pointer to this stub from the thread:
+  // The one that is passed in, is not correct - it points to the code object
+  // that needs to be replaced.
+  __ ldr(CODE_REG,
+         Address(THR, target::Thread::fix_allocation_stub_code_offset()));
+  __ EnterStubFrame();
+  // Preserve type arguments register.
+  __ Push(AllocateObjectABI::kTypeArgumentsReg);
+  // Setup space on stack for return value.
+  __ LoadImmediate(R0, 0);
+  __ Push(R0);
+  __ CallRuntime(kFixAllocationStubTargetRuntimeEntry, 0);
+  // Get Code object result.
+  __ Pop(R0);
+  // Restore type arguments register.
+  __ Push(AllocateObjectABI::kTypeArgumentsReg);
+  // Remove the stub frame.
+  __ LeaveStubFrame();
+  // Jump to the dart function.
+  __ mov(CODE_REG, Operand(R0));
+  __ Branch(FieldAddress(R0, target::Code::entry_point_offset()));
+}
+
 // Input parameters:
 //   R2: smi-tagged argument count, may be zero.
 //   FP[target::frame_layout.param_end_from_fp + 1]: last argument.

runtime/vm/compiler/stub_code_compiler_arm64.cc

@@ -972,6 +972,32 @@ void StubCodeCompiler::GenerateFixAllocationStubTargetStub(
   __ br(R0);
 }
 
+// Called from object allocate instruction when the allocation stub for a
+// generic class has been disabled.
+void StubCodeCompiler::GenerateFixParameterizedAllocationStubTargetStub(
+    Assembler* assembler) {
+  // Load code pointer to this stub from the thread:
+  // The one that is passed in, is not correct - it points to the code object
+  // that needs to be replaced.
+  __ ldr(CODE_REG,
+         Address(THR, target::Thread::fix_allocation_stub_code_offset()));
+  __ EnterStubFrame();
+  // Preserve type arguments register.
+  __ Push(AllocateObjectABI::kTypeArgumentsReg);
+  // Setup space on stack for return value.
+  __ Push(ZR);
+  __ CallRuntime(kFixAllocationStubTargetRuntimeEntry, 0);
+  // Get Code object result.
+  __ Pop(CODE_REG);
+  // Restore type arguments register.
+  __ Pop(AllocateObjectABI::kTypeArgumentsReg);
+  // Remove the stub frame.
+  __ LeaveStubFrame();
+  // Jump to the dart function.
+  __ LoadFieldFromOffset(R0, CODE_REG, target::Code::entry_point_offset());
+  __ br(R0);
+}
+
 // Input parameters:
 //   R2: smi-tagged argument count, may be zero.
 //   FP[target::frame_layout.param_end_from_fp + 1]: last argument.

runtime/vm/compiler/stub_code_compiler_ia32.cc

@@ -558,6 +558,24 @@ void StubCodeCompiler::GenerateFixAllocationStubTargetStub(
   __ int3();
 }
 
+// Called from object allocate instruction when the allocation stub for a
+// generic class has been disabled.
+void StubCodeCompiler::GenerateFixParameterizedAllocationStubTargetStub(
+    Assembler* assembler) {
+  __ EnterStubFrame();
+  // Preserve type arguments register.
+  __ pushl(AllocateObjectABI::kTypeArgumentsReg);
+  __ pushl(Immediate(0));  // Setup space on stack for return value.
+  __ CallRuntime(kFixAllocationStubTargetRuntimeEntry, 0);
+  __ popl(EAX);  // Get Code object.
+  // Restore type arguments register.
+  __ popl(AllocateObjectABI::kTypeArgumentsReg);
+  __ movl(EAX, FieldAddress(EAX, target::Code::entry_point_offset()));
+  __ LeaveFrame();
+  __ jmp(EAX);
+  __ int3();
+}
+
 // Input parameters:
 //   EDX: smi-tagged argument count, may be zero.
 //   EBP[target::frame_layout.param_end_from_fp + 1]: last argument.

runtime/vm/compiler/stub_code_compiler_riscv.cc

@@ -791,6 +791,32 @@ void StubCodeCompiler::GenerateFixAllocationStubTargetStub(
   __ jr(TMP);
 }
 
+// Called from object allocate instruction when the allocation stub for a
+// generic class has been disabled.
+void StubCodeCompiler::GenerateFixParameterizedAllocationStubTargetStub(
+    Assembler* assembler) {
+  // Load code pointer to this stub from the thread:
+  // The one that is passed in, is not correct - it points to the code object
+  // that needs to be replaced.
+  __ lx(CODE_REG,
+        Address(THR, target::Thread::fix_allocation_stub_code_offset()));
+  __ EnterStubFrame();
+  // Preserve type arguments register.
+  __ PushRegister(AllocateObjectABI::kTypeArgumentsReg);
+  // Setup space on stack for return value.
+  __ PushRegister(ZR);
+  __ CallRuntime(kFixAllocationStubTargetRuntimeEntry, 0);
+  // Get Code object result.
+  __ PopRegister(CODE_REG);
+  // Restore type arguments register.
+  __ PopRegister(AllocateObjectABI::kTypeArgumentsReg);
+  // Remove the stub frame.
+  __ LeaveStubFrame();
+  // Jump to the dart function.
+  __ LoadFieldFromOffset(TMP, CODE_REG, target::Code::entry_point_offset());
+  __ jr(TMP);
+}
+
 // Input parameters:
 //   T2: smi-tagged argument count, may be zero.
 //   FP[target::frame_layout.param_end_from_fp + 1]: last argument.

runtime/vm/compiler/stub_code_compiler_x64.cc

@@ -877,6 +877,28 @@ void StubCodeCompiler::GenerateFixAllocationStubTargetStub(
   __ int3();
 }
 
+// Called from object allocate instruction when the allocation stub for a
+// generic class has been disabled.
+void StubCodeCompiler::GenerateFixParameterizedAllocationStubTargetStub(
+    Assembler* assembler) {
+  // Load code pointer to this stub from the thread:
+  // The one that is passed in, is not correct - it points to the code object
+  // that needs to be replaced.
+  __ movq(CODE_REG,
+          Address(THR, target::Thread::fix_allocation_stub_code_offset()));
+  __ EnterStubFrame();
+  // Setup space on stack for return value.
+  __ pushq(AllocateObjectABI::kTypeArgumentsReg);
+  __ pushq(Immediate(0));
+  __ CallRuntime(kFixAllocationStubTargetRuntimeEntry, 0);
+  __ popq(CODE_REG);  // Get Code object.
+  __ popq(AllocateObjectABI::kTypeArgumentsReg);
+  __ movq(RAX, FieldAddress(CODE_REG, target::Code::entry_point_offset()));
+  __ LeaveStubFrame();
+  __ jmp(RAX);
+  __ int3();
+}
+
 // Input parameters:
 //   R10: smi-tagged argument count, may be zero.
 //   RBP[target::frame_layout.param_end_from_fp + 1]: last argument.

runtime/vm/object.cc

@@ -5398,7 +5398,7 @@ void Class::DisableAllocationStub() const {
   }
   ASSERT(!existing_stub.IsDisabled());
   // Change the stub so that the next caller will regenerate the stub.
-  existing_stub.DisableStubCode();
+  existing_stub.DisableStubCode(NumTypeParameters() > 0);
   // Disassociate the existing stub from class.
   untag()->set_allocation_stub(Code::null());
 #endif  // defined(DART_PRECOMPILED_RUNTIME)
@@ -17641,11 +17641,13 @@ void Code::DisableDartCode() const {
                         new_code.UncheckedEntryPointOffset());
 }
 
-void Code::DisableStubCode() const {
+void Code::DisableStubCode(bool is_cls_parameterized) const {
   GcSafepointOperationScope safepoint(Thread::Current());
   ASSERT(IsAllocationStubCode());
   ASSERT(instructions() == active_instructions());
-  const Code& new_code = StubCode::FixAllocationStubTarget();
+  const Code& new_code = is_cls_parameterized
+                             ? StubCode::FixParameterizedAllocationStubTarget()
+                             : StubCode::FixAllocationStubTarget();
   SetActiveInstructions(Instructions::Handle(new_code.instructions()),
                         new_code.UncheckedEntryPointOffset());
 }

runtime/vm/object.h

@@ -6857,7 +6857,7 @@ class Code : public Object {
 
   void DisableDartCode() const;
 
-  void DisableStubCode() const;
+  void DisableStubCode(bool is_cls_parameterized) const;
 
   void Enable() const {
     if (!IsDisabled()) return;

runtime/vm/stub_code_list.h

@@ -75,6 +75,7 @@ namespace dart {
   V(ICCallThroughCode)                                                         \
   V(MegamorphicCall)                                                           \
   V(FixAllocationStubTarget)                                                   \
+  V(FixParameterizedAllocationStubTarget)                                      \
   V(Deoptimize)                                                                \
   V(DeoptimizeLazyFromReturn)                                                  \
   V(DeoptimizeLazyFromThrow)                                                   \