remix-project/libs/remix-astwalker/src/astWalker.ts

import { EventEmitter } from "events";
import { Node, AstNode } from "./index";

export declare interface AstWalker {
  new(): EventEmitter;
}

const isObject = function(obj: any): boolean {
  return obj != null && obj.constructor.name === "Object"
}

export function isAstNode(node: Record<string, unknown>): boolean {
  return (
    isObject(node) &&
    'id' in node &&
    'nodeType' in node &&
    'src' in node
  )
}

export function isYulAstNode(node: Record<string, unknown>): boolean {
  return (
    isObject(node) &&
    'nodeType' in node &&
    'src' in node
  )
}


/**
 * Crawl the given AST through the function walk(ast, callback)
 */
/**
 * visit all the AST nodes
 *
 * @param {Object} ast  - AST node
 * @return EventEmitter
 * event('node', <Node Type | false>) will be fired for every node of type <Node Type>.
 * event('node', "*") will be fired for all other nodes.
 * in each case, if the event emits false it does not descend into children.
 * If no event for the current type, children are visited.
 */
export class AstWalker extends EventEmitter {
  manageCallback(
    node: AstNode,
    callback: Record<string, unknown> | Function // eslint-disable-line @typescript-eslint/ban-types
  ): any {
    // FIXME: we shouldn't be doing this callback determination type on each AST node,
    // since the callback function is set once per walk.
    // Better would be to store the right one as a variable and
    // return that.
    if (node) {
      if ((node).name in callback) {
        return callback[(node).name](node);
      } else {
        return callback["*"](node);
      }
    }
    if (<AstNode>node) {
      if ((<AstNode>node).nodeType in callback) {
        /* istanbul ignore next */
        return callback[(<AstNode>node).nodeType](node);
      } else {
        /* istanbul ignore next */
        return callback["*"](node);
      }
    }
  }

  normalizeNodes(nodes: AstNode[]): AstNode[] {
    // Remove null, undefined and empty elements if any
    nodes = nodes.filter(e => e)

    // If any element in nodes is array, extract its members
    const objNodes = []
    nodes.forEach(x => { 
      if (Array.isArray(x)) objNodes.push(...x)
      else objNodes.push(x)
    });
    
    // Filter duplicate nodes using id field
    const normalizedNodes = []
    objNodes.forEach((element) => {
      const firstIndex = normalizedNodes.findIndex(e => e.id === element.id)
      if(firstIndex == -1) normalizedNodes.push(element)
    })
    return normalizedNodes
  }

  getASTNodeChildren(ast: AstNode): AstNode[] {

    let nodes = ast.nodes             // for ContractDefinition
              || ast.body             // for FunctionDefinition, ModifierDefinition, WhileStatement, DoWhileStatement, ForStatement
              || ast.statements       // for Block, YulBlock
              || ast.members          // for StructDefinition, EnumDefinition
              || ast.overrides        // for OverrideSpecifier
              || ast.parameters       // for ParameterList, EventDefinition
              || ast.declarations     // for VariableDeclarationStatement
              || ast.expression       // for Return, ExpressionStatement, FunctionCall, FunctionCallOptions, MemberAccess
              || ast.components       // for TupleExpression
              || ast.subExpression    // for UnaryOperation
              || ast.eventCall        // for EmitStatement
              || []
    
    // If 'nodes' is not an array, convert it into one, for example: ast.body 
    if(nodes && !Array.isArray(nodes)) {
      const tempArr = []
      tempArr.push(nodes)
      nodes = tempArr
    }
    
    // To break object referencing
    nodes = [...nodes]

    if(ast.nodes && ast.baseContracts?.length) { // for ContractDefinition
        nodes.push(...ast.baseContracts)
    } else if (ast.body && ast.overrides && ast.parameters && ast.returnParameters && ast.modifiers) { // for FunctionDefinition
        nodes.push(ast.overrides)
        nodes.push(ast.parameters)
        nodes.push(ast.returnParameters)
        nodes.push(ast.modifiers)
    } else if(ast.typeName) { // for VariableDeclaration, NewExpression, ElementaryTypeNameExpression
        nodes.push(ast.typeName)
    } else if (ast.body && ast.overrides && ast.parameters) { // for ModifierDefinition
        nodes.push(ast.overrides)
        nodes.push(ast.parameters)
    } else if (ast.modifierName && ast.arguments) { // for ModifierInvocation
        nodes.push(ast.modifierName)
        nodes.push(ast.arguments)
    } else if (ast.parameterTypes && ast.returnParameterTypes) { // for ModifierInvocation
        nodes.push(ast.parameterTypes)
        nodes.push(ast.returnParameterTypes)
    } else if (ast.keyType && ast.valueType) { // for Mapping
        nodes.push(ast.keyType)
        nodes.push(ast.valueType)
    } else if (ast.baseType && ast.length) { // for ArrayTypeName
        nodes.push(ast.baseType)
        nodes.push(ast.length)
    } else if (ast.AST) { // for InlineAssembly
        nodes.push(ast.AST)
    } else if (ast.condition && (ast.trueBody || ast.falseBody || ast.body)) { // for IfStatement, WhileStatement, DoWhileStatement
        nodes.push(ast.condition)
        nodes.push(ast.trueBody)
        nodes.push(ast.falseBody)
    } else if (ast.parameters && ast.block) { // for TryCatchClause
        nodes.push(ast.block)
    } else if (ast.externalCall && ast.clauses) { // for TryStatement
        nodes.push(ast.externalCall)
        nodes.push(ast.clauses)
    } else if (ast.body && ast.condition && ast.initializationExpression && ast.loopExpression) { // for ForStatement
        nodes.push(ast.condition)
        nodes.push(ast.initializationExpression)
        nodes.push(ast.loopExpression)
    } else if (ast.declarations && ast.initialValue) { // for VariableDeclarationStatement
        nodes.push(ast.initialValue)
    } else if (ast.condition && (ast.trueExpression || ast.falseExpression)) { // for Conditional
        nodes.push(ast.condition)
        nodes.push(ast.trueExpression)
        nodes.push(ast.falseExpression)
    } else if (ast.leftHandSide && ast.rightHandSide) { // for Assignment
        nodes.push(ast.leftHandSide)
        nodes.push(ast.rightHandSide)
    } else if (ast.leftExpression && ast.rightExpression) { // for BinaryOperation
        nodes.push(ast.leftExpression)
        nodes.push(ast.rightExpression)
    } else if (ast.expression && (ast.arguments || ast.options)) { // for FunctionCall, FunctionCallOptions
        nodes.push(ast.arguments ? ast.arguments : ast.options)
    } else if (ast.baseExpression && (ast.indexExpression || (ast.startExpression && ast.endExpression))) { // for IndexAccess, IndexRangeAccess
        nodes.push(ast.baseExpression)
        if(ast.indexExpression) nodes.push(ast.indexExpression)
        else {
          nodes.push(ast.startExpression)
          nodes.push(ast.endExpression)
        }
    }
    return this.normalizeNodes(nodes)
  }
  
  // eslint-disable-next-line @typescript-eslint/ban-types, @typescript-eslint/explicit-module-boundary-types
  walk(ast: AstNode, callback?: Function | Record<string, unknown>) {
    if (ast) {
      const children: AstNode[] = this.getASTNodeChildren(ast)
      if (callback) {
        if (callback instanceof Function) {
          callback = Object({ "*": callback });
        }
        if (!("*" in callback)) {
          callback["*"] = function() {
            return true;
          };
        }
        if (this.manageCallback(ast, callback) && children?.length) {
          for (const k in children) {
            const child = children[k];
            this.walk(child, callback);
          }
        }
        } else {
          if (children?.length) {
            for (const k in children) {
              const child = children[k];
              this.emit("node", child);
              this.walk(child);
            }
          }
        }
    }
  }
  // eslint-disable-next-line @typescript-eslint/ban-types, @typescript-eslint/explicit-module-boundary-types
  walkFullInternal(ast: AstNode, callback: Function) {
    if (isAstNode(ast) || isYulAstNode(ast)) {
      // console.log(`XXX id ${ast.id}, nodeType: ${ast.nodeType}, src: ${ast.src}`);
      callback(ast);
      for (const k of Object.keys(ast)) {
        // Possible optimization:
        // if (k in ['id', 'src', 'nodeType']) continue;
        const astItem = ast[k];
        if (Array.isArray(astItem)) {
          for (const child of astItem) {
            if (child) {
              this.walkFullInternal(child, callback);
            }
          }
        } else {
          this.walkFullInternal(astItem, callback);
        }
      }
    }
  }

  // Normalizes parameter callback and calls walkFullInternal
  // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
  walkFull(ast: AstNode, callback: any) {
    if (isAstNode(ast) || isYulAstNode(ast)) return this.walkFullInternal(ast, callback);
  }

  // eslint-disable-next-line @typescript-eslint/ban-types, @typescript-eslint/explicit-module-boundary-types
  walkAstList(sourcesList: Node, cb?: Function) {
    if (cb) {
      if (sourcesList.ast) {
        this.walk(sourcesList.ast, cb);
      }
    } else {
      if (sourcesList.ast) {
        this.walk(sourcesList.ast);
      }
    }
  }
}