Merge next to a preceeding assign

samples/sample1.ts

@@ -1,5 +1,3 @@
-import { sys } from '../types/workflowslib'
-
 function main() {
   const values = [1, 2, 3]
 

src/ast/stepnames.ts

@@ -1,4 +1,5 @@
 import {
+  AssignStepAST,
   ForStepASTNamed,
   NamedWorkflowStep,
   NextStepAST,
@@ -270,13 +271,15 @@ function renameJumpTargets(
   replaceLabels: Map<StepName, StepName>,
 ): WorkflowStepASTWithNamedNested {
   switch (step.tag) {
-    case 'assign':
     case 'call':
     case 'raise':
     case 'return':
     case 'jumptarget':
       return step
 
+    case 'assign':
+      return renameJumpTargetsAssign(step, replaceLabels)
+
     case 'next':
       return renameJumpTargetsNext(step, replaceLabels)
 
@@ -297,6 +300,21 @@ function renameJumpTargets(
   }
 }
 
+function renameJumpTargetsAssign(
+  step: AssignStepAST,
+  replaceLabels: Map<StepName, StepName>,
+): AssignStepAST {
+  if (step.next) {
+    const newLabel = replaceLabels.get(step.next)
+
+    if (newLabel) {
+      return step.withNext(newLabel)
+    }
+  }
+
+  return step
+}
+
 function renameJumpTargetsFor(
   step: ForStepASTNamed,
   replaceLabels: Map<StepName, StepName>,

src/ast/steps.ts

@@ -87,12 +87,26 @@ export class AssignStepAST {
   readonly tag = 'assign'
   readonly assignments: VariableAssignment[]
   readonly label?: string
+  readonly next?: string
 
-  constructor(assignments: VariableAssignment[], label?: string) {
+  constructor(
+    assignments: VariableAssignment[],
+    next?: string,
+    label?: string,
+  ) {
     this.assignments = assignments
+    this.next = next
     this.label = label
   }
 
+  withNext(newNext?: string): AssignStepAST {
+    if (newNext === this.next) {
+      return this
+    } else {
+      return new AssignStepAST(this.assignments, newNext, this.label)
+    }
+  }
+
   withLabel(newLabel?: string): AssignStepAST {
     return new AssignStepAST(this.assignments, newLabel)
   }
@@ -653,6 +667,7 @@ export function renderStep(
         assign: step.assignments.map(([key, val]) => {
           return { [key]: expressionToLiteralValueOrLiteralExpression(val) }
         }),
+        ...(step.next && { next: step.next }),
       }
 
     case 'call':

src/transpiler/statements.ts

@@ -83,11 +83,13 @@ export interface ParsingContext {
 export function parseBlockStatement(
   node: any,
   ctx: ParsingContext,
+  postSteps?: WorkflowStepAST[],
 ): WorkflowStepAST[] {
   assertType(node, BlockStatement)
 
   const body = node.body as any[]
-  return transformAST(body.flatMap((x) => parseStep(x, ctx)))
+  const bodySteps = body.flatMap((x) => parseStep(x, ctx))
+  return transformAST(bodySteps.concat(postSteps ?? []))
 }
 
 function parseStep(node: any, ctx: ParsingContext): WorkflowStepAST[] {
@@ -741,9 +743,8 @@ function whileStatementSteps(
     continueTarget: startOfLoop.label,
     breakTarget: endOfLoop.label,
   })
-
-  const steps: WorkflowStepAST[] = parseBlockStatement(node.body, ctx2)
-  steps.push(new NextStepAST(startOfLoop.label))
+  const postSteps = [new NextStepAST(startOfLoop.label)]
+  const steps = parseBlockStatement(node.body, ctx2, postSteps)
 
   return [
     startOfLoop,

src/transpiler/transformations.ts

@@ -62,10 +62,12 @@ function mergeAssignSteps(steps: WorkflowStepAST[]): WorkflowStepAST[] {
     if (
       current.tag === 'assign' &&
       prev?.tag === 'assign' &&
-      prev.assignments.length < 50
+      prev.assignments.length < 50 &&
+      !prev.next
     ) {
       const merged = new AssignStepAST(
         prev.assignments.concat(current.assignments),
+        current.next,
       )
 
       acc.pop()
@@ -203,7 +205,7 @@ function parseRetryPolicyNumber(
 }
 
 /**
- * Flatten switch conditions that contain a single next step.
+ * Merge a next step to the previous step.
  *
  * For example, transforms this:
  *
@@ -223,10 +225,29 @@ function parseRetryPolicyNumber(
 export function flattenPlainNextConditions(
   steps: WorkflowStepAST[],
 ): WorkflowStepAST[] {
-  return steps.map((step) => (step.tag === 'switch' ? flattenNext(step) : step))
+  return steps.reduce((acc: WorkflowStepAST[], step: WorkflowStepAST) => {
+    if (acc.length > 0 && step.tag === 'next') {
+      // Merge a "next" to the previous "assign" step
+      const prev = acc[acc.length - 1]
+
+      if (prev.tag === 'assign' && !prev.next) {
+        acc.pop()
+        acc.push(prev.withNext(step.target))
+      } else {
+        acc.push(step)
+      }
+    } else if (step.tag === 'switch') {
+      // If the condition steps consists of a single "next", merge it with the condition
+      acc.push(flattenNextToCondition(step))
+    } else {
+      acc.push(step)
+    }
+
+    return acc
+  }, [])
 }
 
-function flattenNext(step: SwitchStepAST): SwitchStepAST {
+function flattenNextToCondition(step: SwitchStepAST): SwitchStepAST {
   const transformedBranches = step.branches.map((cond) => {
     if (!cond.next && cond.steps.length === 1 && cond.steps[0].tag === 'next') {
       const nextStep = cond.steps[0]
@@ -393,7 +414,7 @@ function transformExpressionsAssign(
       newSteps.push(...steps2)
       return [name, ex2] as const
     })
-    newSteps.push(new AssignStepAST(newAssignments, step.label))
+    newSteps.push(new AssignStepAST(newAssignments, step.next, step.label))
     return newSteps
   } else {
     return [step]

test/transpiler.spec.ts

@@ -1191,12 +1191,10 @@ describe('Switch statement', () => {
         - assign2:
             assign:
               - country: Dreamland
-        - next1:
             next: return1
         - assign3:
             assign:
               - country: Bentwood
-        - next2:
             next: return1
         - assign4:
             assign:
@@ -1249,12 +1247,10 @@ describe('Switch statement', () => {
         - assign2:
             assign:
               - country: Dreamland
-        - next1:
             next: end
         - assign3:
             assign:
               - country: Bentwood
-        - next2:
             next: end
         - assign4:
             assign:
@@ -1314,12 +1310,10 @@ describe('Switch statement', () => {
                 - assign2:
                     assign:
                       - country: Dreamland
-                - next1:
                     next: return1
                 - assign3:
                     assign:
                       - country: Bentwood
-                - next2:
                     next: return1
                 - assign4:
                     assign:
@@ -1382,7 +1376,6 @@ describe('Switch statement', () => {
         - assign3:
             assign:
             - country: Dreamland
-        - next1:
             next: return1
         - assign4:
             assign:
@@ -2283,7 +2276,6 @@ describe('Loops', () => {
                   - assign2:
                       assign:
                         - i: \${i - 1}
-                  - next1:
                       next: switch1
     `) as unknown
 
@@ -2430,12 +2422,10 @@ describe('Loops', () => {
                             - assign2:
                                 assign:
                                   - x: \${x * 2}
-                            - next1:
                                 next: switch2
                   - assign3:
                       assign:
                         - x: \${2 * (1 - x)}
-                  - next2:
                       next: switch2
         - return1:
             return: \${x}
@@ -2670,7 +2660,6 @@ describe('Loops', () => {
                   - assign2:
                       assign:
                         - x: \${x * 2}
-                  - next1:
                       next: switch1
         - assign3:
             assign: