Use traitOptions when visiting account nodes

Author: lorisleiva

packages/renderers-rust/src/getRenderMapVisitor.ts

@@ -27,14 +27,14 @@ import {
 import { getTypeManifestVisitor } from './getTypeManifestVisitor';
 import { ImportMap } from './ImportMap';
 import { renderValueNode } from './renderValueNodeVisitor';
-import { getImportFromFactory, LinkOverrides, render, TraitOptions } from './utils';
+import { getImportFromFactory, getTraitsFromNodeFactory, LinkOverrides, render, TraitOptions } from './utils';
 
 export type GetRenderMapOptions = {
     defaultTraitOverrides?: string[];
     dependencyMap?: Record<string, string>;
     linkOverrides?: LinkOverrides;
     renderParentInstructions?: boolean;
-    traitOverrides?: TraitOptions;
+    traitOptions?: TraitOptions;
 };
 
 export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
@@ -44,7 +44,8 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
     const renderParentInstructions = options.renderParentInstructions ?? false;
     const dependencyMap = options.dependencyMap ?? {};
     const getImportFrom = getImportFromFactory(options.linkOverrides ?? {});
-    const typeManifestVisitor = getTypeManifestVisitor({ getImportFrom });
+    const getTraitsFromNode = getTraitsFromNodeFactory(options.traitOptions);
+    const typeManifestVisitor = getTypeManifestVisitor({ getImportFrom, getTraitsFromNode });
 
     return pipe(
         staticVisitor(
@@ -149,6 +150,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                     node.arguments.forEach(argument => {
                         const argumentVisitor = getTypeManifestVisitor({
                             getImportFrom,
+                            getTraitsFromNode,
                             nestedStruct: true,
                             parentName: `${pascalCase(node.name)}InstructionData`,
                         });
@@ -189,6 +191,7 @@ export function getRenderMapVisitor(options: GetRenderMapOptions = {}) {
                     const struct = structTypeNodeFromInstructionArgumentNodes(node.arguments);
                     const structVisitor = getTypeManifestVisitor({
                         getImportFrom,
+                        getTraitsFromNode,
                         parentName: `${pascalCase(node.name)}InstructionData`,
                     });
                     const typeManifest = visit(struct, structVisitor);

packages/renderers-rust/src/getTypeManifestVisitor.ts

@@ -17,7 +17,7 @@ import {
 import { extendVisitor, mergeVisitor, pipe, visit } from '@codama/visitors-core';
 
 import { ImportMap } from './ImportMap';
-import { GetImportFromFunction, rustDocblock } from './utils';
+import { GetImportFromFunction, GetTraitsFromNodeFunction, rustDocblock } from './utils';
 
 export type TypeManifest = {
     imports: ImportMap;
@@ -27,10 +27,11 @@ export type TypeManifest = {
 
 export function getTypeManifestVisitor(options: {
     getImportFrom: GetImportFromFunction;
+    getTraitsFromNode: GetTraitsFromNodeFunction;
     nestedStruct?: boolean;
     parentName?: string | null;
 }) {
-    const { getImportFrom } = options;
+    const { getImportFrom, getTraitsFromNode } = options;
     let parentName: string | null = options.parentName ?? null;
     let nestedStruct: boolean = options.nestedStruct ?? false;
     let inlineStruct: boolean = false;
@@ -50,14 +51,12 @@ export function getTypeManifestVisitor(options: {
                 visitAccount(account, { self }) {
                     parentName = pascalCase(account.name);
                     const manifest = visit(account.data, self);
-                    manifest.imports.add(['borsh::BorshSerialize', 'borsh::BorshDeserialize']);
+                    const traits = getTraitsFromNode(account);
+                    manifest.imports.mergeWith(traits.imports);
                     parentName = null;
                     return {
                         ...manifest,
-                        type:
-                            '#[derive(BorshSerialize, BorshDeserialize, Clone, Debug, Eq, PartialEq)]\n' +
-                            '#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]\n' +
-                            `${manifest.type}`,
+                        type: traits.render + manifest.type,
                     };
                 },