1/**
2 * The Step Functions program compiler: `SfnEffect` AST → a plain JSONata-mode
3 * ASL `definition` object (with embedded `Output`s the engine resolves) plus
4 * the IAM `policyStatements` its Task states require.
5 *
6 * The output is deliberately deterministic — state names and variable names
7 * are per-compile ordinals, no timestamps — so `StateMachine`'s
8 * `normalizeDefinition` drift comparison stays quiet across deploys of an
9 * unchanged program.
10 */
11import * as Data from "effect/Data";
12import type * as Duration from "effect/Duration";
13import type { Input } from "../../../Input.ts";
14import { toSeconds as toWholeSeconds } from "../../../Util/Duration.ts";
15import * as Output from "../../../Output.ts";
16import type { PolicyStatement } from "../../IAM/Policy.ts";
17import {
18 inputExpr,
19 isExpr,
20 nodeOf,
21 renderExprString,
22 renderNode,
23 variableExpr,
24} from "./Jsonata.ts";
25import type { AslNode, RetryOptions } from "./Node.ts";
26import { isSfnEffect, type SfnEffect } from "./Program.ts";
27
28/**
29 * A program failed to compile to ASL — e.g. yielding a non-`Sfn` value
30 * inside `Sfn.gen`, or unreachable steps after an unconditional `Sfn.fail`.
31 */
32export class SfnCompileError extends Data.TaggedError("SfnCompileError")<{
33 readonly message: string;
34}> {}
35
36/** The result of {@link compileProgram}. */
37export interface CompiledProgram {
38 /**
39 * The ASL definition (JSONata query language). A plain object with
40 * embedded `Output`s — exactly what `StateMachineProps.definition`
41 * accepts.
42 */
43 readonly definition: Record<string, unknown>;
44 /**
45 * IAM policy statements collected from `Sfn.invoke` / `Sfn.integrate`
46 * task states, for the state machine's execution role.
47 */
48 readonly policyStatements: Input<PolicyStatement>[];
49}
50
51interface Ctx {
52 readonly nameCounters: Record<string, number>;
53 varCounter: number;
54 readonly policyStatements: Input<PolicyStatement>[];
55 readonly invokedFunctions: Set<string>;
56}
57
58type State = Record<string, any>;
59
60/**
61 * A partially-built subgraph: its entry state, its states, and the names of
62 * states whose `Next`/`End` is still open (patched when the fragment is
63 * chained to a successor or sealed).
64 */
65interface Fragment {
66 readonly startAt: string;
67 readonly states: Record<string, State>;
68 readonly terminals: readonly string[];
69}
70
71const allocName = (ctx: Ctx, prefix: string): string => {
72 const n = ctx.nameCounters[prefix] ?? 0;
73 ctx.nameCounters[prefix] = n + 1;
74 return `${prefix}${n}`;
75};
76
77const allocVar = (ctx: Ctx): string => `v${ctx.varCounter++}`;
78
79/** Chain `next` after `self`, patching `self`'s open terminals. */
80const chain = (self: Fragment, next: Fragment): Fragment => {
81 const states = { ...self.states, ...next.states };
82 for (const terminal of self.terminals) {
83 states[terminal] = { ...states[terminal], Next: next.startAt };
84 }
85 return { startAt: self.startAt, states, terminals: next.terminals };
86};
87
88/** Seal a fragment into a `{ StartAt, States }` machine (`End: true`). */
89const seal = (
90 fragment: Fragment,
91): { StartAt: string; States: Record<string, State> } => {
92 const states = { ...fragment.states };
93 for (const terminal of fragment.terminals) {
94 states[terminal] = { ...states[terminal], End: true };
95 }
96 return { StartAt: fragment.startAt, States: states };
97};
98
99/**
100 * Render a payload/assign value: `Expr` references become `{% ... %}`
101 * JSONata strings, embedded `Output`s pass through untouched (the engine
102 * resolves them before serialization), containers recurse, `undefined`
103 * becomes `null`.
104 */
105const renderValue = (value: unknown): unknown => {
106 if (value === undefined) return null;
107 if (isExpr(value)) return renderExprString(value);
108 if (Output.isOutput(value)) return value;
109 if (Array.isArray(value)) return value.map(renderValue);
110 if (typeof value === "object" && value !== null) {
111 return Object.fromEntries(
112 Object.entries(value).map(([key, entry]) => [key, renderValue(entry)]),
113 );
114 }
115 return value;
116};
117
118/**
119 * Convert a `Duration.Input` to ASL whole seconds (normalizing any
120 * persisted-state `Duration` JSON via the central util). ASL second fields
121 * (`TimeoutSeconds`, `IntervalSeconds`, `Seconds`) must be positive
122 * integers, so sub-second durations round up to 1.
123 */
124export const aslSeconds = (input: Duration.Input): number =>
125 Math.max(1, toWholeSeconds(input)!);
126
127const retryPolicy = (options: RetryOptions): Record<string, unknown> => ({
128 ErrorEquals:
129 options.while === undefined
130 ? ["States.ALL"]
131 : typeof options.while === "string"
132 ? [options.while]
133 : [...options.while],
134 IntervalSeconds: aslSeconds(options.initial ?? "1 second"),
135 BackoffRate: options.backoff ?? 2,
136 MaxAttempts: options.maxAttempts ?? 3,
137 ...(options.maxDelay !== undefined
138 ? { MaxDelaySeconds: aslSeconds(options.maxDelay) }
139 : {}),
140 ...(options.jitter ? { JitterStrategy: "FULL" } : {}),
141});
142
143/**
144 * Does this node compile to exactly one catchable/retryable state
145 * (`Task`/`Parallel`/`Map`)? When true, `retry`/`catch` attach their
146 * `Retry`/`Catch` directly instead of wrapping in a single-branch
147 * `Parallel`.
148 */
149const isSingleCatchableState = (node: AslNode): boolean =>
150 node.kind === "invoke" ||
151 node.kind === "integrate" ||
152 node.kind === "all" ||
153 node.kind === "forEach" ||
154 (node.kind === "retry" && isSingleCatchableState(node.inner.node));
155
156export const compileNode = (ctx: Ctx, node: AslNode, v: string): Fragment => {
157 switch (node.kind) {
158 case "invoke": {
159 const name = allocName(ctx, "Invoke");
160 if (!ctx.invokedFunctions.has(node.fn.LogicalId)) {
161 ctx.invokedFunctions.add(node.fn.LogicalId);
162 ctx.policyStatements.push({
163 Effect: "Allow",
164 Action: ["lambda:InvokeFunction"],
165 Resource: [
166 node.fn.functionArn,
167 // cover qualified invocations (versions/aliases)
168 Output.interpolate`${node.fn.functionArn}:*`,
169 ],
170 } as Input<PolicyStatement>);
171 }
172 return {
173 startAt: name,
174 states: {
175 [name]: {
176 Type: "Task",
177 Resource: "arn:aws:states:::lambda:invoke",
178 Arguments: {
179 FunctionName: node.fn.functionArn,
180 Payload: renderValue(node.payload),
181 },
182 Assign: { [v]: "{% $states.result.Payload %}" },
183 },
184 },
185 terminals: [name],
186 };
187 }
188
189 case "integrate": {
190 const name = allocName(ctx, "Task");
191 ctx.policyStatements.push(
192 ...((node.options.policyStatements ?? []) as Input<PolicyStatement>[]),
193 );
194 return {
195 startAt: name,
196 states: {
197 [name]: {
198 Type: "Task",
199 Resource: node.waitForTaskToken
200 ? `${node.options.resource}.waitForTaskToken`
201 : node.options.resource,
202 ...(node.options.arguments !== undefined
203 ? { Arguments: renderValue(node.options.arguments) }
204 : {}),
205 ...(node.options.timeout !== undefined
206 ? { TimeoutSeconds: aslSeconds(node.options.timeout) }
207 : {}),
208 Assign: { [v]: "{% $states.result %}" },
209 },
210 },
211 terminals: [name],
212 };
213 }
214
215 case "all": {
216 const name = allocName(ctx, "Parallel");
217 return {
218 startAt: name,
219 states: {
220 [name]: {
221 Type: "Parallel",
222 Branches: node.branches.map((branch) => compileBranch(ctx, branch)),
223 Assign: { [v]: "{% $states.result %}" },
224 },
225 },
226 terminals: [name],
227 };
228 }
229
230 case "forEach": {
231 const name = allocName(ctx, "Map");
232 // the item arrives as the ItemProcessor's state input — bind it to a
233 // variable in a prologue Pass so the body can reference it anywhere
234 // ($states.context.Map.Item.Value only exists inside ItemSelector)
235 const itemVar = allocVar(ctx);
236 const body = node.body(variableExpr(itemVar));
237 if (!isSfnEffect(body)) {
238 throw new SfnCompileError({
239 message: "Sfn.forEach body must return an Sfn effect",
240 });
241 }
242 const itemName = allocName(ctx, "Item");
243 const intro: Fragment = {
244 startAt: itemName,
245 states: {
246 [itemName]: {
247 Type: "Pass",
248 Assign: { [itemVar]: "{% $states.input %}" },
249 },
250 },
251 terminals: [itemName],
252 };
253 return {
254 startAt: name,
255 states: {
256 [name]: {
257 Type: "Map",
258 Items: renderExprString(node.items),
259 ...(node.options.concurrency !== undefined
260 ? { MaxConcurrency: node.options.concurrency }
261 : {}),
262 ItemProcessor: {
263 ProcessorConfig: { Mode: "INLINE" },
264 ...compileBranch(ctx, body, intro),
265 },
266 Assign: { [v]: "{% $states.result %}" },
267 },
268 },
269 terminals: [name],
270 };
271 }
272
273 case "sleep": {
274 const name = allocName(ctx, "Wait");
275 return {
276 startAt: name,
277 states: {
278 [name]: {
279 Type: "Wait",
280 Seconds: node.seconds,
281 Assign: { [v]: null },
282 },
283 },
284 terminals: [name],
285 };
286 }
287
288 case "when": {
289 const name = allocName(ctx, "Choice");
290 const onTrue = compileNode(ctx, node.onTrue.node, v);
291 const onFalse = compileNode(
292 ctx,
293 node.onFalse?.node ?? { kind: "succeed", value: null },
294 v,
295 );
296 return {
297 startAt: name,
298 states: {
299 [name]: {
300 Type: "Choice",
301 Choices: [
302 {
303 Condition: renderExprString(node.condition),
304 Next: onTrue.startAt,
305 },
306 ],
307 Default: onFalse.startAt,
308 },
309 ...onTrue.states,
310 ...onFalse.states,
311 },
312 terminals: [...onTrue.terminals, ...onFalse.terminals],
313 };
314 }
315
316 case "match": {
317 const name = allocName(ctx, "Choice");
318 const states: Record<string, State> = {};
319 const terminals: string[] = [];
320 const rules: Record<string, unknown>[] = [];
321 const value = renderNode(nodeOf(node.value));
322 for (const [caseValue, program] of Object.entries(node.cases)) {
323 const fragment = compileNode(ctx, program.node, v);
324 Object.assign(states, fragment.states);
325 terminals.push(...fragment.terminals);
326 rules.push({
327 // string-coerced comparison so numeric/boolean case keys behave
328 Condition: `{% ($string(${value}) = ${JSON.stringify(caseValue)}) %}`,
329 Next: fragment.startAt,
330 });
331 }
332 let defaultTarget: string | undefined;
333 if (node.otherwise !== undefined) {
334 const fragment = compileNode(ctx, node.otherwise.node, v);
335 Object.assign(states, fragment.states);
336 terminals.push(...fragment.terminals);
337 defaultTarget = fragment.startAt;
338 }
339 return {
340 startAt: name,
341 states: {
342 [name]: {
343 Type: "Choice",
344 Choices: rules,
345 ...(defaultTarget !== undefined ? { Default: defaultTarget } : {}),
346 },
347 ...states,
348 },
349 terminals,
350 };
351 }
352
353 case "succeed": {
354 const name = allocName(ctx, "Pass");
355 return {
356 startAt: name,
357 states: {
358 [name]: { Type: "Pass", Assign: { [v]: renderValue(node.value) } },
359 },
360 terminals: [name],
361 };
362 }
363
364 case "fail": {
365 const name = allocName(ctx, "Fail");
366 return {
367 startAt: name,
368 states: {
369 [name]: { Type: "Fail", Error: node.error, Cause: node.cause },
370 },
371 terminals: [],
372 };
373 }
374
375 case "retry": {
376 if (isSingleCatchableState(node.inner.node)) {
377 const fragment = compileNode(ctx, node.inner.node, v);
378 const state = fragment.states[fragment.startAt];
379 state.Retry = [...(state.Retry ?? []), retryPolicy(node.options)];
380 return fragment;
381 }
382 // multi-state program: wrap in a single-branch Parallel so the whole
383 // program re-runs on retry (Effect.retry semantics)
384 const name = allocName(ctx, "Try");
385 return {
386 startAt: name,
387 states: {
388 [name]: {
389 Type: "Parallel",
390 Branches: [compileBranch(ctx, node.inner)],
391 Retry: [retryPolicy(node.options)],
392 Assign: { [v]: "{% $states.result[0] %}" },
393 },
394 },
395 terminals: [name],
396 };
397 }
398
399 case "catch": {
400 const errorVar = allocVar(ctx);
401 let guarded: Fragment;
402 if (isSingleCatchableState(node.inner.node)) {
403 guarded = compileNode(ctx, node.inner.node, v);
404 } else {
405 const name = allocName(ctx, "Try");
406 guarded = {
407 startAt: name,
408 states: {
409 [name]: {
410 Type: "Parallel",
411 Branches: [compileBranch(ctx, node.inner)],
412 Assign: { [v]: "{% $states.result[0] %}" },
413 },
414 },
415 terminals: [name],
416 };
417 }
418 const handlerProgram = node.handler(variableExpr(errorVar));
419 if (!isSfnEffect(handlerProgram)) {
420 throw new SfnCompileError({
421 message: "Sfn.catchTag/catchAll handler must return an Sfn effect",
422 });
423 }
424 const handler = compileNode(ctx, handlerProgram.node, v);
425 const target = guarded.states[guarded.startAt];
426 target.Catch = [
427 ...(target.Catch ?? []),
428 {
429 ErrorEquals: [...node.tags],
430 Assign: { [errorVar]: "{% $states.errorOutput %}" },
431 Next: handler.startAt,
432 },
433 ];
434 return {
435 startAt: guarded.startAt,
436 states: { ...guarded.states, ...handler.states },
437 terminals: [...guarded.terminals, ...handler.terminals],
438 };
439 }
440
441 case "gen": {
442 const iterator = node.body(inputExpr());
443 const fragments: Fragment[] = [];
444 let step = iterator.next();
445 while (!step.done) {
446 const program: unknown = step.value;
447 if (!isSfnEffect(program)) {
448 throw new SfnCompileError({
449 message:
450 "Sfn.gen may only yield Sfn effects (Sfn.invoke, Sfn.when, ...) — real Effects cannot compile to ASL",
451 });
452 }
453 const stepVar = allocVar(ctx);
454 fragments.push(compileNode(ctx, program.node, stepVar));
455 step = iterator.next(variableExpr(stepVar));
456 }
457 const diverged =
458 fragments.length > 0 &&
459 fragments[fragments.length - 1].terminals.length === 0;
460 if (!diverged) {
461 // final Pass assigning the generator's return value
462 const name = allocName(ctx, "Pass");
463 fragments.push({
464 startAt: name,
465 states: {
466 [name]: { Type: "Pass", Assign: { [v]: renderValue(step.value) } },
467 },
468 terminals: [name],
469 });
470 }
471 return fragments.reduce((chained, fragment) => {
472 if (chained.terminals.length === 0) {
473 throw new SfnCompileError({
474 message:
475 "unreachable steps after an unconditional Sfn.fail — move the failure into a branch or make it the final step",
476 });
477 }
478 return chain(chained, fragment);
479 });
480 }
481 }
482};
483
484/**
485 * Compile a program into a sealed `{ StartAt, States }` branch machine
486 * (Parallel branch / Map ItemProcessor). The branch's own result variable is
487 * exposed as the branch output via a terminal Pass state.
488 */
489const compileBranch = (
490 ctx: Ctx,
491 program: SfnEffect<any, any>,
492 intro?: Fragment,
493): { StartAt: string; States: Record<string, State> } => {
494 const branchVar = allocVar(ctx);
495 let fragment = compileNode(ctx, program.node, branchVar);
496 if (intro !== undefined) fragment = chain(intro, fragment);
497 if (fragment.terminals.length === 0) {
498 // the branch always fails — no output state needed
499 return seal(fragment);
500 }
501 const name = allocName(ctx, "Result");
502 return seal(
503 chain(fragment, {
504 startAt: name,
505 states: { [name]: { Type: "Pass", Output: `{% $${branchVar} %}` } },
506 terminals: [name],
507 }),
508 );
509};
510
511/**
512 * Compile a typed Step Functions program to a plain ASL definition object
513 * (JSONata query language, embedded `Output`s intact) plus the IAM policy
514 * statements its task states require. Pure and deterministic; throws
515 * {@link SfnCompileError} on structurally invalid programs.
516 *
517 * This is sugar over the raw path: feed the result straight into the
518 * existing `StateMachine` resource (`StateMachine.fromProgram` does exactly
519 * that), or inspect/extend the definition first — the raw
520 * `definition: Record<string, unknown>` escape hatch stays first-class.
521 */
522export const compileProgram = (
523 program: SfnEffect<any, any>,
524): CompiledProgram => {
525 const ctx: Ctx = {
526 nameCounters: {},
527 varCounter: 0,
528 policyStatements: [],
529 invokedFunctions: new Set(),
530 };
531 const v = allocVar(ctx);
532 const fragment = compileNode(ctx, program.node, v);
533 const full =
534 fragment.terminals.length === 0
535 ? fragment
536 : chain(fragment, {
537 startAt: "Return",
538 states: { Return: { Type: "Pass", Output: `{% $${v} %}` } },
539 terminals: ["Return"],
540 });
541 const machine = seal(full);
542 return {
543 definition: {
544 QueryLanguage: "JSONata",
545 StartAt: machine.StartAt,
546 States: machine.States,
547 },
548 policyStatements: ctx.policyStatements,
549 };
550};