Fix stack height for rewrite or ThreadLocal.withInitial

Bug: b/160484830
Change-Id: I89d10b9d45357b57d32db47737035f38a5cb27f8
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/BackportedMethodRewriter.java b/src/main/java/com/android/tools/r8/ir/desugar/BackportedMethodRewriter.java
index a93e87f..98e481d 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/BackportedMethodRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/BackportedMethodRewriter.java
@@ -1834,6 +1834,7 @@
                   builder -> new ThreadLocalSubclassGenerator(builder, appView));
       eventConsumer.acceptBackportedClass(
           threadLocalSubclass, methodProcessingContext.getMethodContext());
+      localStackAllocator.allocateLocalStack(2);
       return ImmutableList.of(
           new CfNew(threadLocalSubclass.type),
           // Massage the stack with dup_x1 and swap:
diff --git a/src/test/java/com/android/tools/r8/desugar/backports/ThreadLocalBackportTest.java b/src/test/java/com/android/tools/r8/desugar/backports/ThreadLocalBackportTest.java
index 929a3fb..37c87dc 100644
--- a/src/test/java/com/android/tools/r8/desugar/backports/ThreadLocalBackportTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/backports/ThreadLocalBackportTest.java
@@ -4,8 +4,11 @@
 
 package com.android.tools.r8.desugar.backports;
 
+import static com.android.tools.r8.DiagnosticsMatcher.diagnosticMessage;
 import static com.android.tools.r8.utils.AndroidApiLevel.LATEST;
+import static org.hamcrest.CoreMatchers.containsString;
 
+import com.android.tools.r8.TestDiagnosticMessages;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestRunResult;
 import com.android.tools.r8.ToolHelper;
@@ -13,6 +16,7 @@
 import com.android.tools.r8.utils.AndroidApiLevel;
 import com.android.tools.r8.utils.StringUtils;
 import java.util.List;
+import java.util.function.Supplier;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -31,13 +35,23 @@
         getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build());
   }
 
-  private static final String EXPECTED_OUTPUT = StringUtils.lines("Hello, world!");
+  private static final String EXPECTED_OUTPUT = StringUtils.lines("Hello, ", "world!");
+
+  private void checkDiagnostics(TestDiagnosticMessages diagnostics) {
+    if (parameters.isDexRuntime() && parameters.getApiLevel().isLessThan(AndroidApiLevel.N)) {
+      diagnostics.assertWarningsMatch(
+          diagnosticMessage(containsString("Type `java.util.function.Supplier` was not found")));
+    } else {
+      diagnostics.assertNoMessages();
+    }
+  }
 
   @Test
   public void testD8() throws Exception {
     testForD8(parameters.getBackend())
         .addInnerClasses(getClass())
         .setMinApi(parameters.getApiLevel())
+        .compileWithExpectedDiagnostics(this::checkDiagnostics)
         .run(parameters.getRuntime(), TestClass.class)
         .apply(this::checkExpected);
   }
@@ -80,9 +94,14 @@
 
   static class TestClass {
 
+    public static ThreadLocal<String> createThreadLocalWithInitial(Supplier<String> supplier) {
+      return ThreadLocal.withInitial(supplier);
+    }
+
     public static void main(String[] args) {
-      ThreadLocal<String> threadLocal = ThreadLocal.withInitial(() -> "Hello, world!");
+      ThreadLocal<String> threadLocal = ThreadLocal.withInitial(() -> "Hello, ");
       System.out.println(threadLocal.get());
+      System.out.println(createThreadLocalWithInitial(() -> "world!").get());
     }
   }
 }