运行时反馈式TreeShaking

618 阅读3分钟

tfjs3.0更新的新treeshaking功能, 详细了解可探索tfjs-treeshaking-test, 分两步, 一步是通过tf.profile收集运行时使用到的kernel, 第二部, 把使用到的kernel, 填入custom_tfjs.json执行cli工具把对应kernel的文件单独导出

image.png

image.png

所以从上面可以稍微抽象什么是运行时反馈式treeshaking: 2个步骤

  1. 第一步记录运行时代码执行日志
  2. 第二步根据执行日志情况, 把未执行的部分去除, 只保留执行过的.

设计

接下来设计一个rollup插件: hotcode, 所以需要思考一下这些问题

  1. 如何记录执行日志?
  2. 那些代码块需要记录?
  3. 如何删除未使用代码?

先回答那些代码块需要被记录: 函数体if else switch分支, 因为这些包含的代码都是可能被执行.

然后可以通过给这些代码块注入一个辅助函数调用来记录执行次数, 比如函数体注入一个__HOTCODE__.fire(blockId), 通过一个blockId表示这个函数体, 这样就能统计执行次数了.

再通过MagicString根据blockId的触发情况删除对应的代码块

magicString.overwrite(node.start, node.end, '{}', overwriteCfg);

然后需要设计插件API, slot是注入记录代码块执行次数阶段, remove是根据log删除无用代码块阶段

export interface HotCodeProps {
  log?: any; // 运行时
  mode?: 'slot' | 'remove' | 'debug';
  removeSwitch?: boolean; // switch case的remove方式 可能对pass throught的写法有影响, 开个开关,当有问题可手动关闭
  sourceMap?: boolean;
  include?: string | RegExp | ReadonlyArray<string | RegExp> | null;
  exclude?: string | RegExp | ReadonlyArray<string | RegExp> | null;
}

实现

blockId由文件路径长度8的md5 + 一个自增序号组成, 因为rollup文件处理顺序是并非固定, 我们还需要注入一个模块

import { walk } from 'estree-walker';
import MagicString from 'magic-string';
import { AcornNode, Plugin } from 'rollup';
import { createHash } from 'crypto';
import { createFilter } from '@rollup/pluginutils';

export interface HotCodeProps {
  log?: any;
  mode?: 'slot' | 'remove' | 'debug';
  removeSwitch?: boolean;
  sourceMap?: boolean;
  include?: string | RegExp | ReadonlyArray<string | RegExp> | null;
  exclude?: string | RegExp | ReadonlyArray<string | RegExp> | null;
}

function getMd5(str) {
  const hash = createHash('sha256');
  hash.update(str);
  return hash.digest('hex');
}

const overwriteCfg = { storeName: true };

export default function hotCode(options: HotCodeProps = {}): Plugin {
  const { sourceMap = false, log = null, mode, removeSwitch = true } = options;
  if (mode !== 'slot' && mode !== 'remove' && mode !== 'debug') return null;

  const isSlotMode = mode === 'slot';
  const isRemoveMode = mode === 'remove';
  const isDebugMode = mode === 'debug';
  const filter = createFilter(options.include, options.exclude);
  const unusedOnRuntime = id => log && log[id] === undefined;

  function getInjectCode(blockId) {
    const errorCode =
      isDebugMode && unusedOnRuntime(blockId) ? `__HOTCODE__.error('${blockId}');` : '';
    return `__HOTCODE__.fire('${blockId}');${errorCode}`;
  }

  return {
    name: 'hotCode',
    transform(code, filePath) {
      if (!filter(filePath)) return null;
      if (filePath.endsWith('hot-code-mod.js')) return null;

      let ast: AcornNode = null;
      try {
        ast = this.parse(code);
      } catch (err) {
        this.warn({
          code: 'PARSE_ERROR',
          message: `rollup-plugin-hotCode: failed to parse ${filePath}.`,
        });
      }
      if (!ast) return null;

      let fileBlockId = 0;
      const fileId = getMd5(filePath).slice(0, 8);
      const getBlockId = () => `${fileId}-${fileBlockId++}`;

      const magicString = new MagicString(code);
      // 注入辅助模块
      magicString.prepend(`import __HOTCODE__ from '${__dirname}/hot-code-mod';\n`);

      // 注入括号的方式, 适用于代码块(BlockStatement)
      function injectWithBrackets(node, removeFN?: (node: any) => void) {
        if (node && node.start !== undefined && node.start != node.end) {
          const blockId = getBlockId();
          const fireCode = getInjectCode(blockId);

          if (node.type === 'BlockStatement') {
            if (isSlotMode || isDebugMode) magicString.prependRight(node.start + 1, fireCode);
            else if (isRemoveMode && unusedOnRuntime(blockId)) {
              magicString.overwrite(node.start, node.end, '{}', overwriteCfg);
            }
          // 特殊处理变量声明不能使用{}包装
          } else if (node.type === 'VariableDeclaration') {
            if (isSlotMode || isDebugMode) {
              magicString.prependLeft(node.start, fireCode);
            } else if (isRemoveMode && unusedOnRuntime(blockId)) {
              if (removeFN) removeFN(node);
            }
          // 默认使用{}包装, 不影响执行结果
          } else {
            if (isSlotMode || isDebugMode) {
              magicString.prependLeft(node.start, '{' + fireCode);
              magicString.prependRight(node.end, '}');
            } else if (isRemoveMode && unusedOnRuntime(blockId))
              if (removeFN) removeFN(node);
              else magicString.overwrite(node.start, node.end, '{}', overwriteCfg);
          }
        }
      }

      // 注入分号的方式, 适用于表达式
      function injectWithCommas(node, removeFN?: (node: any) => void) {
        if (node && node.start !== undefined && node.start != node.end) {
          const blockId = getBlockId();
          const fireCode = getInjectCode(blockId).replace(/;/g, ',');

          if (isSlotMode || isDebugMode) {
            magicString.prependLeft(node.start, '(' + fireCode);
            magicString.prependRight(node.end, ')');
          } else if (isRemoveMode && unusedOnRuntime(blockId)) {
            if (removeFN) removeFN(node);
            else magicString.overwrite(node.start, node.end, '0', overwriteCfg);
          }
        }
      }

      walk(ast, {
        leave(node, parent) {
          if (sourceMap && node.start !== undefined) {
            magicString.addSourcemapLocation(node.start);
            magicString.addSourcemapLocation(node.end);
          }

          switch (node.type) {
            case 'ArrowFunctionExpression':
              node.expression ? injectWithCommas(node.body) : injectWithBrackets(node.body);
              break;
            case 'FunctionDeclaration':
            case 'FunctionExpression':
              // 两种避免IIFE未被rollup treeshaking处理方式, 一种是不处理, 一种是返回空对象
              if (!(parent && parent.type === 'CallExpression' && parent.callee === node))
                injectWithBrackets(node.body);
              break;
            case 'IfStatement':
              injectWithBrackets(node.consequent);
              injectWithBrackets(node.alternate);
              break;
            case 'SwitchStatement':
              if (removeSwitch)
                node.cases.forEach(i => {
                  injectWithBrackets(i.consequent[0], node => {
                    magicString.prependLeft(node.start, 'break;');
                  });
                });
              break;
          }
        },
      });

      return {
        code: magicString.toString(),
        map: sourceMap ? magicString.generateMap({ hires: true }) : null,
      };
    },
  };
}

辅助模块hot-code-mod.ts

const fireLog = {};
const fireError = new Set();

console.fireLog = fireLog;
console.fireError = fireError;
// 微信小游戏console不能扩展...
setInterval(() => { console.log(fireLog); }, 3000);

export function fire(id: number) {
  if (fireLog[id] == undefined) fireLog[id] = 0;
  fireLog[id]++;
}

export function error(id: number) {
  fireError.add(id);
}

效果

oasis的FlappyBird移植微信小程序版为例子, 使用platformize轻松实现, 多小程序的代码复用

先构造对照, 无运行时反馈式treeshaking, 可以看到包体积为545kb

> pnpm build --filter platformize-oasis-wechat-game

image.png

然后开始第一步: 插入代码执行记录

> pnpm build-hotcode-slot --filter platformize-oasis-wechat-game

然后触发游戏逻辑, 然后再记录代码块执行情况

image.png

image.png

image.png

然后执行第二步: 根据执行记录移除未使用的代码块

pnpm build-hotcode-remove --filter platformize-oasis-wechat-game

image.png

可以看到小程序体积从545kb变到了290kb, 体积减少了46.7%, 并且FlappyBird也可以正常游戏

同样也测试了platformizethree-wechat-simple, 小程序包大小606kb, 运行时反馈式treeshaking后只有321kb, 体积减少47.0%

目前还属于prototype阶段, 包含hotcode plugin的版本尚未发布, 感兴趣的同学可根据仓库文档跑起来试试, 也欢迎提交PR/issue👏🏻

platformize 是一个支持把 js 库中浏览器 api 改用定制 polyfill 的构建插件, 并提供特定库的定制适配, 比如threejs@0.133.0 / oasis@0.6.3 / playcanvas@1.50.0 / pixi@6.2.1

image.png

单测

最后补充单测用例

function printTips() {
  console.log('a');
  tips.forEach((tip, i) => console.log(`Tip ${i}:` + tip));
}

const a = {
  b() {
    console.log('b');
  },
};

if (true) {
  console.log('c');
} else {
  console.log('d');
}

switch (2) {
  case '2':
    console.log('e');
    console.log('e1');
    break;
  case '3':
    const a = 0;
    console.log('e');
    console.log('e1');
    break;
  case '4': {
    const a = 0;
    console.log('e');
    console.log('e1');
    break;
  }
  default:
    console.log('f');
    console.log('f1');
    break;
}

1 ? 3 : 2;

if (1) console.log('g');

const fn = () => console.log('h');
const fn1 = () => {
  console.log('h');
};

class A {
  m0() {
    console.log('i');
  }
  m1() {
    console.log('j');
    [0, 1].forEach(i => {
      console.log('k');
    });
    const m2 = () => {
      console.log('l');
    };

    if (2) {
      console.log('m');
      m2();
    } else {
      console.log('n');
    }
  }
}

const f = () => (console.log('o'), Date.now);
const f1 = () => Date.now;
const f2 = () => Date.now;

const fixRollupTreeShaking = function () {
  function a() {}
  return a;
};

fixRollupTreeShaking._temp = 0;

单测执行结果快照

# Snapshot report for `test/test-hot-code.js`

The actual snapshot is saved in `test-hot-code.js.snap`.

Generated by [AVA](https://avajs.dev).

## hotcode

> Snapshot 1

    `import __HOTCODE__ from '/Users/bytedance/opensource/platformize/packages/platformize/dist-plugin/hot-code-mod';␊
    function printTips() {__HOTCODE__.fire('047765e8-1');␊
      console.log('a');␊
      tips.forEach((tip, i) => (__HOTCODE__.fire('047765e8-0'),console.log(\`Tip ${i}:\` + tip)));␊
    }␊
    ␊
    const a = {␊
      b() {__HOTCODE__.fire('047765e8-2');␊
        console.log('b');␊
      },␊
    };␊
    ␊
    if (true) {__HOTCODE__.fire('047765e8-3');␊
      console.log('c');␊
    } else {__HOTCODE__.fire('047765e8-4');␊
      console.log('d');␊
    }␊
    ␊
    switch (2) {␊
      case '2':␊
        {__HOTCODE__.fire('047765e8-5');console.log('e');}␊
        console.log('e1');␊
        break;␊
      case '3':␊
        __HOTCODE__.fire('047765e8-6');const a = 0;␊
        console.log('e');␊
        console.log('e1');␊
        break;␊
      case '4': {__HOTCODE__.fire('047765e8-7');␊
        const a = 0;␊
        console.log('e');␊
        console.log('e1');␊
        break;␊
      }␊
      default:␊
        {__HOTCODE__.fire('047765e8-8');console.log('f');}␊
        console.log('f1');␊
        break;␊
    }␊
    ␊
    1 ? 3 : 2;␊
    ␊
    if (1) {__HOTCODE__.fire('047765e8-9');console.log('g');}␊
    ␊
    const fn = () => (__HOTCODE__.fire('047765e8-10'),console.log('h'));␊
    const fn1 = () => {__HOTCODE__.fire('047765e8-11');␊
      console.log('h');␊
    };␊
    ␊
    class A {␊
      m0() {__HOTCODE__.fire('047765e8-12');␊
        console.log('i');␊
      }␊
      m1() {__HOTCODE__.fire('047765e8-17');␊
        console.log('j');␊
        [0, 1].forEach(i => {__HOTCODE__.fire('047765e8-13');␊
          console.log('k');␊
        });␊
        const m2 = () => {__HOTCODE__.fire('047765e8-14');␊
          console.log('l');␊
        };␊
    ␊
        if (2) {__HOTCODE__.fire('047765e8-15');␊
          console.log('m');␊
          m2();␊
        } else {__HOTCODE__.fire('047765e8-16');␊
          console.log('n');␊
        }␊
      }␊
    }␊
    ␊
    const f = () => ((__HOTCODE__.fire('047765e8-18'),console.log('o'), Date.now));␊
    const f1 = () => (__HOTCODE__.fire('047765e8-19'),Date.now);␊
    const f2 = () => (__HOTCODE__.fire('047765e8-20'),Date.now);␊
    ␊
    const fixRollupTreeShaking = function () {__HOTCODE__.fire('047765e8-22');␊
      function a() {__HOTCODE__.fire('047765e8-21');}␊
      return a;␊
    };␊
    ␊
    fixRollupTreeShaking._temp = 0;␊
    `