Function getKernelMappingForFile

tfjs-converter/scripts/kernels_to_ops.ts:71–127

(source: SourceFile)

};

function getKernelMappingForFile(source: SourceFile) {
  const switchStatement = getSwitchStatement(source);
  if (switchStatement == null) {
    throw new Error('No switch statment found in executor');
  }
  const caseClauses = switchStatement.getClauses();

  const kernelsToOp: KernelMapping = {};
  let currentClauseGroup: string[] = [];

  // Loop through clauses until you reach one that has a block or return.
  // This allows us to coalesce fallthrough case blocks in a switch statement.
  caseClauses.forEach((caseClause: CaseOrDefaultClause) => {
    if (caseClause instanceof CaseClause) {
      let kernelName;
      caseClause.forEachChild(clausePart => {
        const kind = clausePart.getKindName();
        if (kind === 'StringLiteral') {
          kernelName = clausePart.getText().replace(/\'/g, '');
          currentClauseGroup.push(kernelName);
        }
        if (kind === 'Block' || kind === 'ReturnStatement') {
          // We have reached a code block, all the previously captured
          // kernels use this block as their execution path.

          // Parse the code block and determing all the tfc.*() function calls
          // used.
          const callExprs =
              clausePart.getDescendantsOfKind(SyntaxKind.CallExpression);
          const tfOpsCallExprs =
              callExprs.filter(expr => expr.getText().match(/ops/));
          const tfSymbols: Set<string> = new Set();
          for (const tfOpsCall of tfOpsCallExprs) {
            const tfOpsCallStr = tfOpsCall.getText();
            const functionCallMatcher = /(ops\.([\w\.]*)\()/g;
            const matches = tfOpsCallStr.match(functionCallMatcher);
            if (matches != null && matches.length > 0) {
              for (const match of matches) {
                // extract the method name (and any namespaces used to call it)
                const symbolMatcher = /(ops\.([\w\.]*)\()/;
                const symbol = match.match(symbolMatcher)[2];
                tfSymbols.add(symbol);
              }
            }
          }
          for (const kern of currentClauseGroup) {
            kernelsToOp[kern] = Array.from(tfSymbols);
          }
          // Reset the clause tracker as we are moving to a new set of kernels
          currentClauseGroup = [];
        }
      });
    }
  });

  return kernelsToOp;
}