Create a BranchCtx for Raster Pipeline branch ops.

Previously, these ops just pointed directly at an int*. Now they
point at a BranchCtx* (which contains an int). At present, this
should result in the same data in the Raster Pipeline's alloc; an
int doesn't behave differently than a struct holding an int.

One minor tweak is that we now put the branch target contexts
into the alloc as branch instructions are emitted, rather than
creating them in a slab at the start of the alloc. This might
improve locality _very slightly_ but in practice I assume it's
immaterial.

These changes will make it easier to add new types of RP branching
ops that use more context data than just a single int*.

Change-Id: I41321743fbdf2c61f33cd28b3c3b29cf974ba558
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/633577
Auto-Submit: John Stiles <johnstiles@google.com>
Reviewed-by: Brian Osman <brianosman@google.com>
Commit-Queue: John Stiles <johnstiles@google.com>
diff --git a/src/core/SkRasterPipelineOpContexts.h b/src/core/SkRasterPipelineOpContexts.h
index 08b0cac..7f8b914 100644
--- a/src/core/SkRasterPipelineOpContexts.h
+++ b/src/core/SkRasterPipelineOpContexts.h
@@ -146,4 +146,8 @@
     uint16_t offsets[16];  // values must be byte offsets (4 * highp-stride * component-index)
 };
 
+struct SkRasterPipeline_BranchCtx {
+    int offset;
+};
+
 #endif  // SkRasterPipelineOpContexts_DEFINED
diff --git a/src/opts/SkRasterPipeline_opts.h b/src/opts/SkRasterPipeline_opts.h
index 899916a..515f45b 100644
--- a/src/opts/SkRasterPipeline_opts.h
+++ b/src/opts/SkRasterPipeline_opts.h
@@ -3176,16 +3176,16 @@
     update_execution_mask();
 }
 
-STAGE_BRANCH(branch_if_any_active_lanes, int* offset) {
-    return any(execution_mask()) ? *offset : 1;
+STAGE_BRANCH(branch_if_any_active_lanes, SkRasterPipeline_BranchCtx* ctx) {
+    return any(execution_mask()) ? ctx->offset : 1;
 }
 
-STAGE_BRANCH(branch_if_no_active_lanes, int* offset) {
-    return any(execution_mask()) ? 1 : *offset;
+STAGE_BRANCH(branch_if_no_active_lanes, SkRasterPipeline_BranchCtx* ctx) {
+    return any(execution_mask()) ? 1 : ctx->offset;
 }
 
-STAGE_BRANCH(jump, int* offset) {
-    return *offset;
+STAGE_BRANCH(jump, SkRasterPipeline_BranchCtx* ctx) {
+    return ctx->offset;
 }
 
 STAGE_TAIL(immediate_f, void* ctx) {
diff --git a/src/sksl/codegen/SkSLRasterPipelineBuilder.cpp b/src/sksl/codegen/SkSLRasterPipelineBuilder.cpp
index 3b6cb8e..bf80af5 100644
--- a/src/sksl/codegen/SkSLRasterPipelineBuilder.cpp
+++ b/src/sksl/codegen/SkSLRasterPipelineBuilder.cpp
@@ -820,9 +820,10 @@
     int currentStack = 0;
     int mostRecentRewind = 0;
 
-    // Allocate buffers for branch targets (used when running the program) and labels (only needed
-    // during initial program construction).
-    int* branchTargets = alloc->makeArrayDefault<int>(fNumBranches);
+    // Allocate buffers for branch targets and labels; these are needed during initial program
+    // construction, to convert labels into actual offsets into the pipeline and fix up branches.
+    SkTArray<SkRasterPipeline_BranchCtx*> branchTargets;
+    branchTargets.push_back_n(fNumBranches, (SkRasterPipeline_BranchCtx*)nullptr);
     SkTArray<int> labelOffsets;
     labelOffsets.push_back_n(fNumLabels, -1);
     SkTArray<int> branchGoesToLabel;
@@ -859,7 +860,7 @@
 
             case BuilderOp::jump:
             case BuilderOp::branch_if_any_active_lanes:
-            case BuilderOp::branch_if_no_active_lanes:
+            case BuilderOp::branch_if_no_active_lanes: {
                 // If we have already encountered the label associated with this branch, this is a
                 // backwards branch. Add a stack-rewind immediately before the branch to ensure that
                 // long-running loops don't use an unbounded amount of stack space.
@@ -873,12 +874,15 @@
                 // targets at the end and fix them up.
                 SkASSERT(inst.fImmA >= 0 && inst.fImmA < fNumLabels);
                 SkASSERT(currentBranchOp >= 0 && currentBranchOp < fNumBranches);
-                branchTargets[currentBranchOp] = pipeline->size();
+                SkRasterPipeline_BranchCtx* branchCtx = alloc->make<SkRasterPipeline_BranchCtx>();
+                branchCtx->offset = pipeline->size();
+
+                branchTargets[currentBranchOp] = branchCtx;
                 branchGoesToLabel[currentBranchOp] = inst.fImmA;
-                pipeline->push_back({(RPOp)inst.fOp, &branchTargets[currentBranchOp]});
+                pipeline->push_back({(RPOp)inst.fOp, branchCtx});
                 ++currentBranchOp;
                 break;
-
+            }
             case BuilderOp::init_lane_masks:
                 pipeline->push_back({RPOp::init_lane_masks, nullptr});
                 break;
@@ -1132,9 +1136,9 @@
 
     // Fix up every branch target.
     for (int index = 0; index < fNumBranches; ++index) {
-        int branchFromIdx = branchTargets[index];
+        int branchFromIdx = branchTargets[index]->offset;
         int branchToIdx = labelOffsets[branchGoesToLabel[index]];
-        branchTargets[index] = branchToIdx - branchFromIdx;
+        branchTargets[index]->offset = branchToIdx - branchFromIdx;
     }
 }
 
diff --git a/tests/SkRasterPipelineTest.cpp b/tests/SkRasterPipelineTest.cpp
index 92c3628..4d6bb60 100644
--- a/tests/SkRasterPipelineTest.cpp
+++ b/tests/SkRasterPipelineTest.cpp
@@ -1555,7 +1555,8 @@
 
     alignas(64) static constexpr float kColorDarkRed[4] = {0.5f, 0.0f, 0.0f, 0.75f};
     alignas(64) static constexpr float kColorGreen[4]   = {0.0f, 1.0f, 0.0f, 1.0f};
-    const int offset = 2;
+    SkRasterPipeline_BranchCtx ctx;
+    ctx.offset = 2;
 
     // An array of all zeros.
     alignas(64) static constexpr int32_t kNoLanesActive[4 * SkRasterPipeline_kMaxStride_highp] = {};
@@ -1567,16 +1568,16 @@
     // Make a program which conditionally branches past two append_constant_color ops.
     SkArenaAlloc alloc(/*firstHeapAllocation=*/256);
     SkRasterPipeline p(&alloc);
-    p.append_constant_color(&alloc, kColorDarkRed);                  // set the color to dark red
+    p.append_constant_color(&alloc, kColorDarkRed);                    // set the color to dark red
     p.append(SkRasterPipelineOp::load_dst, kNoLanesActive);            // make no lanes active
-    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &offset); // do not skip past next line
-    p.append_constant_color(&alloc, kColorGreen);                    // set the color to green
+    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &ctx);    // do not skip past next line
+    p.append_constant_color(&alloc, kColorGreen);                      // set the color to green
     p.append(SkRasterPipelineOp::load_dst, oneLaneActive);             // set one lane active
-    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &offset); // skip past next line
-    p.append_constant_color(&alloc, kColorDarkRed);                  // (not executed)
+    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &ctx);    // skip past next line
+    p.append_constant_color(&alloc, kColorDarkRed);                    // (not executed)
     p.append(SkRasterPipelineOp::init_lane_masks);                     // set all lanes active
-    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &offset); // skip past next line
-    p.append_constant_color(&alloc, kColorDarkRed);                  // (not executed)
+    p.append(SkRasterPipelineOp::branch_if_any_active_lanes, &ctx);    // skip past next line
+    p.append_constant_color(&alloc, kColorDarkRed);                    // (not executed)
     p.append(SkRasterPipelineOp::store_src, slots);                    // store final color
     p.run(0,0,1,1);
 
@@ -1598,7 +1599,8 @@
     alignas(64) static constexpr float kColorBlack[4]   = {0.0f, 0.0f, 0.0f, 0.0f};
     alignas(64) static constexpr float kColorRed[4]     = {1.0f, 0.0f, 0.0f, 1.0f};
     alignas(64) static constexpr float kColorBlue[4]    = {0.0f, 0.0f, 1.0f, 1.0f};
-    const int offset = 2;
+    SkRasterPipeline_BranchCtx ctx;
+    ctx.offset = 2;
 
     // An array of all zeros.
     alignas(64) static constexpr int32_t kNoLanesActive[4 * SkRasterPipeline_kMaxStride_highp] = {};
@@ -1610,16 +1612,16 @@
     // Make a program which conditionally branches past a append_constant_color op.
     SkArenaAlloc alloc(/*firstHeapAllocation=*/256);
     SkRasterPipeline p(&alloc);
-    p.append_constant_color(&alloc, kColorBlack);                    // set the color to black
+    p.append_constant_color(&alloc, kColorBlack);                      // set the color to black
     p.append(SkRasterPipelineOp::init_lane_masks);                     // set all lanes active
-    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &offset);  // do not skip past next line
-    p.append_constant_color(&alloc, kColorRed);                      // sets the color to red
+    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &ctx);     // do not skip past next line
+    p.append_constant_color(&alloc, kColorRed);                        // sets the color to red
     p.append(SkRasterPipelineOp::load_dst, oneLaneActive);             // set one lane active
-    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &offset);  // do not skip past next line
+    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &ctx);     // do not skip past next line
     p.append(SkRasterPipelineOp::swap_rb);                             // swap R and B (making blue)
     p.append(SkRasterPipelineOp::load_dst, kNoLanesActive);            // make no lanes active
-    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &offset);  // skip past next line
-    p.append_constant_color(&alloc, kColorBlack);                    // (not executed)
+    p.append(SkRasterPipelineOp::branch_if_no_active_lanes, &ctx);     // skip past next line
+    p.append_constant_color(&alloc, kColorBlack);                      // (not executed)
     p.append(SkRasterPipelineOp::store_src, slots);                    // store final blue color
     p.run(0,0,1,1);