fix: support synthetic prop mutation, and duplicate mutations on the …

packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react.test.ts.snap

@@ -5261,6 +5261,115 @@ export default function ColorChangeOnClick(
 }
 `;
 
+exports[`amplify render tests mutations supports multiple actions pointing to the same value 1`] = `
+Object {
+  "componentText": "/* eslint-disable */
+import React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+  useStateMutationAction,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { Button, Flex, FlexProps, Text } from \\"@aws-amplify/ui-react\\";
+
+export type ButtonsToggleStateProps = React.PropsWithChildren<
+  Partial<FlexProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function ButtonsToggleState(
+  props: ButtonsToggleStateProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  const [fooBarValueChildren, setFooBarValueChildren] =
+    useStateMutationAction(\\"Baz\\");
+  const fooButtonClick = () => {
+    setFooBarValueChildren(\\"Foo\\");
+  };
+  const barButtonClick = () => {
+    setFooBarValueChildren(\\"Bar\\");
+  };
+  return (
+    /* @ts-ignore: TS2322 */
+    <Flex {...rest} {...getOverrideProps(overrides, \\"ButtonsToggleState\\")}>
+      <Text
+        children={fooBarValueChildren}
+        {...getOverrideProps(overrides, \\"FooBarValue\\")}
+      ></Text>
+      <Button
+        children=\\"Set to Foo\\"
+        onClick={() => {
+          fooButtonClick();
+        }}
+        {...getOverrideProps(overrides, \\"FooButton\\")}
+      ></Button>
+      <Button
+        children=\\"Set to Bar\\"
+        onClick={() => {
+          barButtonClick();
+        }}
+        {...getOverrideProps(overrides, \\"BarButton\\")}
+      ></Button>
+    </Flex>
+  );
+}
+",
+  "declaration": undefined,
+  "renderComponentToFilesystem": [Function],
+}
+`;
+
+exports[`amplify render tests mutations supports mutations on synthetic props 1`] = `
+Object {
+  "componentText": "/* eslint-disable */
+import React from \\"react\\";
+import {
+  EscapeHatchProps,
+  getOverrideProps,
+  useStateMutationAction,
+} from \\"@aws-amplify/ui-react/internal\\";
+import { Button, Flex, FlexProps, Text } from \\"@aws-amplify/ui-react\\";
+
+export type MutationWithSyntheticPropProps = React.PropsWithChildren<
+  Partial<FlexProps> & {
+    overrides?: EscapeHatchProps | undefined | null;
+  }
+>;
+export default function MutationWithSyntheticProp(
+  props: MutationWithSyntheticPropProps
+): React.ReactElement {
+  const { overrides, ...rest } = props;
+  const [fooBarValueChildren, setFooBarValueChildren] =
+    useStateMutationAction(\\"Baz\\");
+  const fooButtonClick = () => {
+    setFooBarValueChildren(\\"Foo\\");
+  };
+  return (
+    /* @ts-ignore: TS2322 */
+    <Flex
+      {...rest}
+      {...getOverrideProps(overrides, \\"MutationWithSyntheticProp\\")}
+    >
+      <Text
+        children={fooBarValueChildren}
+        {...getOverrideProps(overrides, \\"FooBarValue\\")}
+      ></Text>
+      <Button
+        children=\\"Set to Foo\\"
+        onClick={() => {
+          fooButtonClick();
+        }}
+        {...getOverrideProps(overrides, \\"FooButton\\")}
+      ></Button>
+    </Flex>
+  );
+}
+",
+  "declaration": undefined,
+  "renderComponentToFilesystem": [Function],
+}
+`;
+
 exports[`amplify render tests primitives Built-in Iconset 1`] = `
 "/* eslint-disable */
 import React from \\"react\\";

packages/codegen-ui-react/lib/__tests__/studio-ui-codegen-react.test.ts

@@ -337,6 +337,14 @@ describe('amplify render tests', () => {
     it('internal mutation', () => {
       expect(generateWithAmplifyRenderer('workflow/internalMutation')).toMatchSnapshot();
     });
+
+    it('supports mutations on synthetic props', () => {
+      expect(generateWithAmplifyRenderer('workflow/mutationWithSyntheticProp')).toMatchSnapshot();
+    });
+
+    it('supports multiple actions pointing to the same value', () => {
+      expect(generateWithAmplifyRenderer('workflow/buttonsToggleState')).toMatchSnapshot();
+    });
   });
 
   describe('default value', () => {

packages/codegen-ui-react/lib/__tests__/workflow/mutation.test.ts

@@ -43,7 +43,7 @@ describe('getComponentStateReferences', () => {
       children: [
         {
           componentType: 'TextField',
-          name: 'UsernameTextField',
+          name: 'UserNameTextField',
           properties: {
             label: {
               value: 'Username',

packages/codegen-ui-react/lib/__tests__/workflow/utils.test.ts

@@ -0,0 +1,40 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+ */
+
+import { getChildPropMappingForComponentName } from '../../workflow/utils';
+
+describe('getChildPropMappingForComponentName', () => {
+  const componentNameToTypeMap = {
+    MyFlex: 'Flex',
+    MyButton: 'Button',
+  };
+
+  test('throws on missing name mapping', () => {
+    expect(() => {
+      getChildPropMappingForComponentName(componentNameToTypeMap, 'MissingComponent');
+    }).toThrowErrorMatchingInlineSnapshot(
+      `"Invalid definition, found reference to component name MissingComponent which wasn't found in the schema."`,
+    );
+  });
+
+  test('returns mapping for synthetic mapped prop', () => {
+    expect(getChildPropMappingForComponentName(componentNameToTypeMap, 'MyButton')).toEqual('label');
+  });
+
+  test('returns undefined for non-mapped component', () => {
+    expect(getChildPropMappingForComponentName(componentNameToTypeMap, 'MyFlex')).toBeUndefined();
+  });
+});

packages/codegen-ui-react/lib/react-studio-template-renderer.ts

@@ -29,9 +29,9 @@ import {
   StudioComponentVariant,
   StudioComponentSimplePropertyBinding,
   handleCodegenErrors,
-  StudioComponentChild,
   StateStudioComponentProperty,
   MutationActionSetStateParameter,
+  buildComponentNameToTypeMap,
 } from '@aws-amplify/codegen-ui';
 import { EOL } from 'os';
 import ts, {
@@ -1028,7 +1028,7 @@ export abstract class ReactStudioTemplateRenderer extends StudioTemplateRenderer
     const actions = getComponentActions(component);
     if (actions) {
       return actions.map(({ action, identifier }) =>
-        buildUseActionStatement(action, identifier, this.importCollection),
+        buildUseActionStatement(component, action, identifier, this.importCollection),
       );
     }
     return [];
@@ -1117,10 +1117,12 @@ export abstract class ReactStudioTemplateRenderer extends StudioTemplateRenderer
       return;
     }
 
+    const componentNameToTypeMap = buildComponentNameToTypeMap(this.component);
+
     this.component.variants.forEach((variant) => {
       Object.entries(variant.overrides).forEach(([name, value]) => {
         const propsInOverrides = value;
-        const componentType = this.getComponentTypeFromName(name);
+        const componentType = componentNameToTypeMap[name];
         if (componentType && isPrimitive(componentType)) {
           const childrenPropMapping = PrimitiveChildrenPropMapping[Primitive[componentType as Primitive]];
           if (childrenPropMapping !== undefined) {
@@ -1135,39 +1137,5 @@ export abstract class ReactStudioTemplateRenderer extends StudioTemplateRenderer
     });
   }
 
-  /**
-   * Build a mapping from component name to component type,
-   * pull out the specific component type for the given name.
-   * This is required for synthetic prop mapping of variants,
-   * since we're performing it at the top level, but need
-   * to know the underlying type given the override key.
-   */
-  private getComponentTypeFromName(name: string): string | undefined {
-    const componentNameToTypeMap = this.buildComponentNameToTypeMap(this.component);
-    return componentNameToTypeMap[name];
-  }
-
-  /**
-   * Helper to recurse through the component tree and build the name to type mapping.
-   */
-  private buildComponentNameToTypeMap(component: StudioComponent | StudioComponentChild): Record<string, string> {
-    const localMap: Record<string, string> = {};
-    if (component.name) {
-      localMap[component.name] = component.componentType;
-    }
-    if (component.children) {
-      Object.entries(
-        component.children
-          .map((child) => this.buildComponentNameToTypeMap(child))
-          .reduce((previous, next) => {
-            return { ...previous, ...next };
-          }, {}),
-      ).forEach(([name, componentType]) => {
-        localMap[name] = componentType;
-      });
-    }
-    return localMap;
-  }
-
   abstract renderJsx(component: StudioComponent): JsxElement | JsxFragment | JsxSelfClosingElement;
 }

packages/codegen-ui-react/lib/workflow/action.ts

@@ -22,9 +22,11 @@ import {
   ActionStudioComponentEvent,
   StudioComponentProperty,
   MutationAction,
+  buildComponentNameToTypeMap,
 } from '@aws-amplify/codegen-ui';
 import { isActionEvent, propertyToExpression, getSetStateName } from '../react-component-render-helper';
 import { ImportCollection, ImportSource, ImportValue } from '../imports';
+import { getChildPropMappingForComponentName } from './utils';
 
 enum Action {
   'Amplify.Navigate' = 'Amplify.Navigate',
@@ -87,12 +89,13 @@ export function getActionIdentifier(componentName: string | undefined, event: st
 }
 
 export function buildUseActionStatement(
+  component: StudioComponent,
   action: ActionStudioComponentEvent,
   identifier: string,
   importCollection: ImportCollection,
 ): Statement {
   if (isMutationAction(action)) {
-    return buildMutationActionStatement(action, identifier);
+    return buildMutationActionStatement(component, action, identifier);
   }
 
   const actionHookImportValue = getActionHookImportValue(action.action);
@@ -115,8 +118,16 @@ export function buildUseActionStatement(
   );
 }
 
-export function buildMutationActionStatement(action: MutationAction, identifier: string) {
-  const setState = getSetStateName(action.parameters.state);
+export function buildMutationActionStatement(component: StudioComponent, action: MutationAction, identifier: string) {
+  const componentNameToTypeMap = buildComponentNameToTypeMap(component);
+  const { componentName, property } = action.parameters.state;
+  const childrenPropMapping = getChildPropMappingForComponentName(componentNameToTypeMap, componentName);
+  const stateReference =
+    childrenPropMapping !== undefined && property === childrenPropMapping
+      ? { ...action.parameters.state, property: 'children' }
+      : action.parameters.state;
+
+  const setState = getSetStateName(stateReference);
 
   return factory.createVariableStatement(
     undefined,

packages/codegen-ui-react/lib/workflow/mutation.ts

@@ -21,6 +21,7 @@ import {
   ActionStudioComponentEvent,
   StateReference,
   StudioGenericEvent,
+  buildComponentNameToTypeMap,
 } from '@aws-amplify/codegen-ui';
 import {
   isActionEvent,
@@ -32,8 +33,16 @@ import {
 } from '../react-component-render-helper';
 import { ImportCollection, ImportValue } from '../imports';
 import { mapGenericEventToReact } from './events';
+import { getChildPropMappingForComponentName } from './utils';
 
-export function getComponentStateReferences(component: StudioComponent | StudioComponentChild) {
+export function getComponentStateReferences(component: StudioComponent) {
+  const stateReferences = getComponentStateReferencesHelper(component);
+  const componentNameToTypeMap = buildComponentNameToTypeMap(component);
+  const mappedStateReferences = mapSyntheticReferences(stateReferences, componentNameToTypeMap);
+  return mappedStateReferences;
+}
+
+function getComponentStateReferencesHelper(component: StudioComponent | StudioComponentChild) {
   const stateReferences: StateReference[] = [];
 
   if (component.properties) {
@@ -54,15 +63,27 @@ export function getComponentStateReferences(component: StudioComponent | StudioC
 
   if (component.children) {
     component.children.forEach((child) => {
-      stateReferences.push(...getComponentStateReferences(child));
+      stateReferences.push(...getComponentStateReferencesHelper(child));
     });
   }
 
-  // TODO: dedupe state references
-
   return stateReferences;
 }
 
+function mapSyntheticReferences(
+  stateReferences: StateReference[],
+  componentNameToTypeMap: Record<string, string>,
+): StateReference[] {
+  return stateReferences.map((stateReference) => {
+    const { componentName, property } = stateReference;
+    const childrenPropMapping = getChildPropMappingForComponentName(componentNameToTypeMap, componentName);
+    if (childrenPropMapping !== undefined && property === childrenPropMapping) {
+      return { ...stateReference, property: 'children' };
+    }
+    return stateReference;
+  });
+}
+
 export function getActionStateParameters(action: ActionStudioComponentEvent): StateStudioComponentProperty[] {
   if (action.parameters) {
     return Object.entries(action.parameters)
@@ -120,6 +141,23 @@ export function buildOpeningElementControlEvents(stateName: string, event: strin
   );
 }
 
+/**
+ * Dedupes state references by componentName + property, returning a consolidate
+ * list, stripping the `set` property from them. We do this by serializing to json,
+ * collecting in a set, then deserializing.
+ */
+function dedupeStateReferences(stateReferences: StateReference[]): StateReference[] {
+  const dedupedSerializedRefs = [
+    ...new Set(
+      stateReferences.map((stateReference) => {
+        const { componentName, property } = stateReference;
+        return JSON.stringify({ componentName, property });
+      }),
+    ),
+  ];
+  return dedupedSerializedRefs.map((ref) => JSON.parse(ref));
+}
+
 export function buildStateStatements(
   component: StudioComponent,
   stateReferences: StateReference[],
@@ -128,7 +166,8 @@ export function buildStateStatements(
   if (stateReferences.length > 0) {
     importCollection.addMappedImport(ImportValue.USE_STATE_MUTATION_ACTION);
   }
-  return stateReferences.map((stateReference) => {
+  const dedupedStateReferences = dedupeStateReferences(stateReferences);
+  return dedupedStateReferences.map((stateReference) => {
     return factory.createVariableStatement(
       undefined,
       factory.createVariableDeclarationList(