Rudimentary Boxart plugin transforming bablyon tree to flatten functions

babel-boxart-plugin-rudimentary.js

// function gen() {
//   return (function(fn) {
//     const f = function(t, state, begin, end) {
//       var b = {b: 1};
//       var c = {f2: function(t) {return t;}};
//       state.left = c.f2(b.b, state, begin, end) * 2;
//       return state;
//     };
//     return f;
//   })((function(fn) {
//     const f = function(t) {return fn(t * 2) * 2;};
//     return f;
//   })((function(fn) {
//     const f = function(t) {return fn(t * 2) * 2;};
//     return f;
//   })((function(fn) {
//     const f = function(t) {return fn(t * 2) * 2;};
//     return f;
//   })((function(c) {
//     const f = function(t) {return t * c;};
//     return f;
//   })(3)))));
// }


export default function(babel) {
  const { types: t, traverse } = babel;

  let values = {};

  const memberLookup = (path, state) => {
    let stack = [];
    let node = path.node;
    while (t.isMemberExpression(node)) {
      if (node.computed) {return;}
      if (!t.isIdentifier(node.property)) {return;}
      stack.unshift(node.property.name);
      node = node.object;
    }
    if (!t.isIdentifier(node)) {return;}
    let value = state[node.name];
    while (t.isObjectExpression(value) && stack.length > 0) {
      const name = stack.shift();
      value = value.properties.find(prop => prop.key.name === name);
      if (value) {
        value = value.value;
      }
    }
    if (t.isLiteral(value) || t.isIdentifier(value) || t.isFunctionExpression(value)) {
      path.replaceWith(t.cloneDeep(value));
    }
  };

  const inlineFunctions = {
    FunctionExpression(path, state) {},
    Identifier(path, state) {
      // if (
      //   !path.parentPath.isVariableDeclarator() ||
      //   path.parent.id !== path.node
      // ) {
      //   if (t.isLiteral(state[path.node.name]) || t.isFunction(state[path.node.name])) {
      //     path.replaceWith(t.cloneDeep(state[path.node.name]));
      //   }
      // }
    },
    VariableDeclarator: {
      enter(path, state) {
        path.scope.rename(path.node.id.name);
      },
      exit(path, state) {
        if (path.get("init").isIdentifier() && state[path.node.id.name]) {
          path.get('init').replaceWith(t.cloneDeep(state[path.node.id.name]));
        }
        if (path.get("init").isFunctionExpression()) {
          state[path.node.id.name] = path.node.init;
        }
        if (path.get("init").isNumericLiteral()) {
          state[path.node.id.name] = path.node.init;
        }
        if (path.get("init").isObjectExpression()) {
          state[path.node.id.name] = path.node.init;
        }
      }
    },
    CallExpression: {
      enter(path, state) {
        if (path.get("callee").isMemberExpression()) {
          memberLookup(path.get('callee'), state);
          // path.get("callee").replaceWith(t.cloneDeep(state[path.node.callee.name]));
        }
        if (
          path.get("callee").isIdentifier() &&
          state[path.node.callee.name]
        ) {
          path.get("callee").replaceWith(t.cloneDeep(state[path.node.callee.name]));
        }
        if (path.get("callee").isFunctionExpression()) {
          state = Object.assign({}, state, { __parentstate: state });
          const args = path.node.arguments;
          const params = path.node.callee.params;
          params.forEach((p, i) => {
          });
        }
      },
      exit(path, state) {
        if (path.get("callee").isFunctionExpression()) {
          const args = path.node.arguments;
          const params = path.node.callee.params;
          params.forEach((p, i) => {
            path.get("callee").scope.rename(p.name);
            path
              .getStatementParent()
              .insertBefore(
                t.variableDeclaration(
                  "const",
                  [t.variableDeclarator(t.identifier(p.name), args[i])]
                )
              );
            if (
              path.get(`arguments.${i}`).isIdentifier() &&
              t.isFunctionExpression(state[args[i].name])
            ) {
              state[p.name] = state[args[i].name];
            }
            if (path.get(`arguments.${i}`).isLiteral()) {
              state[p.name] = args[i];
            }
          });
        }
        if (path.get("callee").isFunctionExpression()) {
          path.get("callee.body").node.body.forEach(expr => {
            if (t.isReturnStatement(expr)) {
              path.replaceWith(t.cloneDeep(expr.argument));
            } else {
              path.getStatementParent().insertBefore(t.cloneDeep(expr));
            }
          });
        }
      }
    }
  };

  const lookupAndOps = {
    VariableDeclarator: {
      exit(path, state) {
        if (
          path.get('init').isLiteral() ||
          path.get('init').isObjectExpression()
        ) {
          state[path.node.id.name] = path.node.init;
        }
      },
    },
    MemberExpression(path, state) {
      let stack = [];
      let node = path.node;
      while (t.isMemberExpression(node)) {
        if (node.computed) {return;}
        if (!t.isIdentifier(node.property)) {return;}
        stack.unshift(node.property.name);
        node = node.object;
      }
      // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${String(stack)}`)));
      if (!t.isIdentifier(node)) {return;}
      let value = state[node.name];
      // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${JSON.stringify(node.name)}`)));
      // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${JSON.stringify(value)}`)));
      while (t.isObjectExpression(value) && stack.length > 0) {
        const name = stack.shift();
        // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${name}`)));
        value = value.properties.find(prop => prop.key.name === name);
        if (value) {
          value = value.value;
        }
      }
      // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${stack.length} ${JSON.stringify(value)}`)));
      if (t.isLiteral(value)) {
        path.replaceWith(t.cloneDeep(value));
      }
    },
    Identifier(path, state) {
      if (
        path.parent.id !== path.node
      ) {
        if (state[path.node.name]) {
          path.replaceWith(t.cloneDeep(state[path.node.name]));
        }
      }
    },
    BinaryExpression: {
      exit(path, state) {
        // const e = path.evaluate();
        // if (e.confident) {
        // }
        if (
          path.get('left').isLiteral() && path.get('right').isLiteral() &&
          path.node.operator === '*'
        ) {
          const e = path.evaluate();
          if (e.confident && typeof e.value === 'number') {
            path.replaceWith(t.numericLiteral(e.value));
          }
        }
        if (
          path.get('left').isBinaryExpression() &&
          path.node.left.operator === '*' && path.node.operator === '*' &&
          path.get('left.right').isLiteral() && path.get('right').isLiteral()
        ) {
          path.replaceWith(
            t.binaryExpression(
              '*',
              path.node.left.left,
              t.binaryExpression('*', path.node.left.right, path.node.right)
            )
          );
          const e = path.get('right').evaluate();
          if (e.confident && typeof e.value === 'number') {
            path.get('right').replaceWith(t.numericLiteral(e.value));
          }
        }
      },
    },
  };

  const refCount = {
    Identifier(path, state) {
      if (path.getStatementParent().parentPath.isProgram()) {
        return;
      }
      if (path.parent.id === path.node) {
        if (!(state[path.node.name] || {}).node) {
          // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(path.node.name)));
          state[path.node.name] = {node: path.node, refs: [], refsFrom: []};
        }
      }
      if (
        path.getStatementParent().isVariableDeclaration() &&
        path.findParent(t.isVariableDeclarator).node.id.name !== path.node.name
      ) {
        const declId = path.findParent(t.isVariableDeclarator).node.id.name;
        const id = path.node.name;
        if ((state[id] || {}).node) {
          // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(String([declId, id]))));
          if (!(state[declId] || {}).node) {
            state[declId] = {node: path.parent.id, refs: [], refsFrom: []};
          }
          state[declId].refs.push(id);
          if (!(state[id] || {}).node) {
            state[id] = {node: path.node, refs: [], refsFrom: []};
          }
          state[id].refsFrom.push(declId);
        }
      }
      if (path.getStatementParent().isReturnStatement()) {
        const id = path.node.name;
        if (!(state[id] || {}).node) {
          state[id] = {node: path.node, refs: [], refsFrom: []};
        }
        state[id].refsFrom.push('__return__');
      }
    },
  };

  const deadCode = {
    Identifier(path, state) {
      if (
        path.parent.id === path.node
      ) {
        // path.getStatementParent().insertBefore(t.expressionStatement(t.stringLiteral(`${path.node.name}: ${String((state[path.node.name] || {refsFrom: []}).refsFrom.length)}`)));
      }
      if (
        path.parent.id === path.node &&
        state[path.node.name] &&
        state[path.node.name].refsFrom.length === 0
      ) {
        if (!path.getStatementParent().parentPath.isProgram()) {
          path.getStatementParent().remove();
        }
      }
    },
  };

  return {
    visitor: {
      Program(path) {
        path.traverse(inlineFunctions, {});
        path.traverse(lookupAndOps, {});
        let refs = {};
        path.traverse(refCount, refs);
        let lastRefs = refs;
        let removed = Infinity;
        while (removed > 0) {
          removed = 0;
          refs = {};
          path.traverse(deadCode, lastRefs);
          path.traverse(refCount, refs);
          if (Object.keys(refs).length < Object.keys(lastRefs).length) {
            removed = 1;
          }
          lastRefs = refs;
        }
        path.skip();
      },
    }
  };
}