Link external rust IDL structs into a main IDL. Created for sniper.xyz sniper market IDL linking.

anchor_dependency_linker.ts

/* eslint-disable no-prototype-builtins */
import { Project, SourceFile, VariableDeclarationKind } from 'ts-morph';
import prettier from 'prettier';
import * as fs from 'fs';
import colors from 'colors';

//  This works if you reference the struct in the instruction parameters EX:
//  pub fn buy_tensor_pool_listing_v1<'info>(
//     ctx: Context<'_, '_, '_, 'info, BuyTensorPoolListingV1<'info>>,
//     config: tensorswap::PoolConfig,
//  )
//
//  In the IDL file the struct will be name "tensorswap::PoolConfig"
//
//  This script would replace the struct name to "TensorswapPoolConfig" and add that struct to the IDL file
//
//  This script is run after program is built and before the tests are run.

const sourceIdlJson = './target/idl/sniper_market.json';
const sourceIdlTs = './target/types/sniper_market.ts';
const externalIdlJson = './cpi-autogen/tensorswap/src/idl.json';

const PROGRAM_NAME = 'SniperMarket';
const EXTERNAL_PROGRAM_NAME = 'Tensorswap';
const EXTERNAL_PROGRAM_NAME_PREFIX = 'tensorswap::';

(async () => {
  console.log(
    colors.green(
      `Bundling ${sourceIdlTs} IDL to include external types from ${EXTERNAL_PROGRAM_NAME}...`,
    ),
  );

  const sourceIdl: { types: any[]; accounts: any[] } = JSON.parse(
    fs.readFileSync(sourceIdlJson, 'utf-8'),
  );
  const tensorswapIdl = JSON.parse(fs.readFileSync(externalIdlJson, 'utf-8'));

  const sniperMarketTypeDefs: { name: string }[] = sourceIdl.types;
  const tensorTypesDefs: { name: string; type: { kind: 'struct'; fields: Object[] } }[] =
    tensorswapIdl.types;

  // When you reference a type from cpi rust package, it will be prefixed with the package name in IDL.
  const tensorRefStructsFromSniperMarketIDL = findValuesAndReplaceWithPrefix(
    sourceIdl,
    EXTERNAL_PROGRAM_NAME_PREFIX, // 'tensorswap::'
    EXTERNAL_PROGRAM_NAME, // 'Tensorswap'
  );

  const tensorReferencedTypes = findReferencedStructs(
    tensorRefStructsFromSniperMarketIDL,
    tensorTypesDefs,
    true,
  ).filter((v, i, a) => a.findIndex((t) => t.name === v.name) === i);

  const renamedTensorReferencedTypes = tensorReferencedTypes.map((typeDef) =>
    addPrefixToNamesAndDefinedValues(EXTERNAL_PROGRAM_NAME, typeDef),
  );

  sourceIdl.types = [...sniperMarketTypeDefs, ...renamedTensorReferencedTypes];

  // throw if there are still duplicate types
  const duplicateTypes = sourceIdl.accounts.filter(
    (v, i, a) => a.findIndex((t) => t.name === v.name) !== i,
  );

  if (duplicateTypes.length > 0) {
    throw new Error(`Duplicate types found: ${duplicateTypes.map((t) => t.name).join(', ')}`);
  }

  const project = new Project();
  const sourceFile = project.createSourceFile('sniper_market.ts', '', { overwrite: true });

  // Take json alteration and generates modified IDL typescript type
  generateType(sourceFile, sourceIdl, PROGRAM_NAME) + ';';

  // Take json alteration and create a const variable of the IDL
  sourceFile.addVariableStatement({
    declarationKind: VariableDeclarationKind.Const,
    isExported: true,
    declarations: [
      {
        name: 'IDL',
        type: PROGRAM_NAME,
        initializer: JSON.stringify(sourceIdl, null, 2),
      },
    ],
  });

  // Save to a file
  const formattedType = prettier.format(sourceFile.getFullText(), { parser: 'typescript' });
  fs.writeFileSync(sourceIdlTs, formattedType);
  fs.writeFileSync(sourceIdlJson, JSON.stringify(sourceIdl, null, 2));

  console.log(colors.green(`Bundling complete...`));
})();

function generateType(sourceFile: SourceFile, obj: Object, typeName: string) {
  function generateTypeNode(obj) {
    if (Array.isArray(obj)) {
      const arrayElementTypes = obj.map((element) => generateTypeNode(element)).join(', ');
      return `[${arrayElementTypes}]`;
    } else if (typeof obj === 'object') {
      const properties: string[] = [];
      for (const key in obj) {
        const type = generateTypeNode(obj[key]);
        properties.push(`${key}: ${type}`);
      }
      return `{ ${properties.join(', ')} }`;
    } else if (typeof obj === 'string') {
      const escapedValue = obj.replace(/'/g, "\\'");
      return `'${escapedValue}'`;
    } else if (typeof obj === 'number') {
      return `${obj}`;
    } else if (typeof obj === 'boolean') {
      return `${obj}`;
    } else {
      return typeof obj;
    }
  }
  const typeNode = generateTypeNode(obj);
  sourceFile.addTypeAlias({ name: typeName, type: typeNode, isExported: true });

  return typeNode;
}

function findValuesAndReplaceWithPrefix(
  obj: { [key: string]: any },
  prefix: string,
  replace: string = '',
) {
  let result: string[] = [];

  for (const key in obj) {
    const value = obj[key];

    if (typeof value === 'string' && value.startsWith(prefix)) {
      result.push(value);

      obj[key] = value.replace(prefix, replace);
    } else if (Array.isArray(value)) {
      value.forEach((item) => {
        result = result.concat(findValuesAndReplaceWithPrefix(item, prefix, replace));
      });
    } else if (typeof value === 'object' && value !== null) {
      result = result.concat(findValuesAndReplaceWithPrefix(value, prefix, replace));
    }
  }

  return result;
}

export type IDLType = { name: string; type: { kind: 'struct'; fields: Object[] } };

function findReferencedStructs(
  depthStructs: string[],
  typeDefs: IDLType[],
  isFirstDepth: boolean = false,
) {
  const structsToReplace: { name: string; type: { kind: 'struct'; fields: Object[] } }[] = [];

  for (const depthStruct of depthStructs) {
    const name = isFirstDepth ? depthStruct.split('::')[1] : depthStruct;
    const tensorStruct = typeDefs.find((def) => def.name === name);

    if (!tensorStruct) {
      continue;
    }

    const depthDefinedStructs = findDefinedValues(tensorStruct.type);
    const nextDepthStructs = findReferencedStructs(depthDefinedStructs, typeDefs);

    structsToReplace.push(tensorStruct, ...nextDepthStructs);
  }

  return structsToReplace;
}

function addPrefixToNamesAndDefinedValues(prefix: string, object: IDLType) {
  if (object.hasOwnProperty('name')) {
    object.name = `${prefix}${object.name}`;
  }

  Object.values(object.type).forEach((value) => {
    findAndReplaceDefinedValues(value, (str) => `${prefix}${str}`);
  });

  return object;
}

function findDefinedValues(obj) {
  const result: string[] = [];

  function search(object: Object) {
    for (const key in object) {
      if (key === 'defined') {
        result.push(object[key] as string);
      } else if (typeof object[key] === 'object') {
        search(object[key] as Object);
      } else if (typeof object[key] === 'string') {
        // do nothing
      } else if (typeof object[key] === 'number') {
        // do nothing
      } else {
        console.error(`not supported ${key} ${object[key]} ${JSON.stringify(obj)}`);
        throw new Error(`not supported ${key} ${object[key]} ${JSON.stringify(obj)}`);
      }
    }
  }

  search(obj);
  return result;
}

function findAndReplaceDefinedValues(obj: IDLType | Object, replace: (str: string) => string) {
  function searchAndReplace(object: Object) {
    for (const key in object) {
      if (key === 'defined') {
        object[key] = replace(object[key] as string);
      } else if (typeof object[key] === 'object') {
        searchAndReplace(object[key] as Object);
      } else if (Array.isArray(object[key])) {
        object[key] = object[key].map((item) => searchAndReplace(item));
      } else if (typeof object[key] === 'string') {
        // do nothing
      } else if (typeof object[key] === 'number') {
        // do nothing
      } else {
        console.error(`not supported ${key} ${object[key]} ${JSON.stringify(obj)}`);
        throw new Error(`not supported ${key} ${object[key]} ${JSON.stringify(obj)}`);
      }
    }
  }

  searchAndReplace(obj);
}