feat(aws-stepfunctions): a method to retrieve all possible next states from a given initial state

Author: andrestone

packages/@aws-cdk/aws-stepfunctions/README.md

@@ -728,6 +728,14 @@ new stepfunctions.Parallel(this, 'All jobs')
     .branch(new MyJob(this, 'Slow', { jobFlavor: 'slow' }).prefixStates());
 ```
 
+To retrieve all possible subsequent states from an initial state, you can use the
+method `State.findReachableStates`:
+
+```typescript
+        import * as sf from '@aws-cdk/aws-stepfunctions';
+        const nextState = sf.State.findReachableStates(startState); // All reachable / next states
+```
+
 ## Activity
 
 **Activities** represent work that is done on some non-Lambda worker pool. The

packages/@aws-cdk/aws-stepfunctions/lib/states/state.ts

@@ -73,10 +73,28 @@ export abstract class State extends cdk.Construct implements IChainable {
     }
   }
 
+  /**
+   * Find the set of states reachable through transitions from the given start state
+   */
+  public static findReachableStates(start: State, options: FindStateOptions = {}): State[] {
+    const visited = new Set<State>();
+    const ret = new Set<State>();
+    const queue = [start];
+    while (queue.length > 0) {
+      const state = queue.splice(0, 1)[0]!;
+      if (visited.has(state)) { continue; }
+      visited.add(state);
+      const outgoing = state.outgoingTransitions(options);
+      queue.push(...outgoing);
+      ret.add(state);
+    }
+    return Array.from(ret);
+  }
+
   /**
    * Find the set of end states states reachable through transitions from the given start state
    */
-  public static findReachableEndStates(start: State, options: FindStateOptions = {}) {
+  public static findReachableEndStates(start: State, options: FindStateOptions = {}): State[] {
     const visited = new Set<State>();
     const ret = new Set<State>();
     const queue = [start];

packages/@aws-cdk/aws-stepfunctions/test/states-language.test.ts

@@ -614,6 +614,26 @@ describe('States Language', () => {
     expect(() => new stepfunctions.Parallel(stack, 'Parallel')
       .branch(state1.next(state2))
       .branch(state2)).toThrow();
+  }),
+
+  test('Can retrieve possible states from initial state', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const state1 = new stepfunctions.Pass(stack, 'State1');
+    const state2 = new stepfunctions.Pass(stack, 'State2');
+    const state3 = new stepfunctions.Pass(stack, 'State3');
+
+    const definition = state1
+      .next(state2)
+      .next(state3);
+
+    // WHEN
+    const states = stepfunctions.State.findReachableStates(definition.startState);
+
+    // THEN
+    expect(state1.toStateJson()).toStrictEqual(states[0].toStateJson());
+    expect(state2.toStateJson()).toStrictEqual(states[1].toStateJson());
+    expect(state3.toStateJson()).toStrictEqual(states[2].toStateJson());
   });
 });