Node.js でつくる WASMコンパイラー - 07：ユーザ定義関数を実装する

はじめに

Node.jsで小さなプログラミング言語を作ってみるシリーズを、「ミニインタープリター」「ミニコンパイラー」とやってきました。そして三部作(?)の最後として、 ミニNode.jsからWASMを生成する小さなコンパイラーに取り組んでいます。

• Node.jsでつくるNode.js-WASMコンパイラ - もくじ
• 前回の記事 ... 06：文字列出力を実装しFizzBuzzを実現する

これまでの取り組み

• Node.jsでつくるNode.js（ミニインタープリター） - もくじ
• Node.jsでつくるNode.jsミニコンパイラ - もくじ

今回の実現したいこと

これまで関数はあらかじめ用意している2つの組み込み関数だけが利用できました。今回は自由に関数を定義して、呼び出せる様にします。

WATでの関数の定義

JSで次の様な関数を考えます。

function add(x, y) {
  return x + y;
}

対応するWATでは、次の様に関数を記述します。

(func $add (param $x i32) (param $y i32) (result i32)
  get_local $x
  get_local $y
  i32.add
  return
)

WATでの関数の呼び出し

すでに putn() / puts() で見た様に、関数呼び出しは次の様に引数をスタックに積んでから呼び出します。

  i32.const 1
  i32.const 2
  call $add

このようなWATを生成できるようにして行きます。

ユーザ定義関数の実装

関数の定義

ユーザ定義関数が登場した場合、パーサーは次の単純化ASTを生成します。

• [ 'func_def', 関数名, [ 引数の配列 ], [ 関数の中身 ] ]

上記の add() 関数の場合は次の様になります。

[ 'func_def',
    'add',
    [ 'x', 'y' ],
    [ 'ret', 
        [ '+', [ 'var_ref', 'x' ], [ 'var_ref', 'y' ] ] 
    ] 
]

これを扱えるように、コンパイラーを拡張します。

function generate(tree, indent, gctx, lctx) {
  // ... 省略 ...

  // --- func_def user_defined function ---
  if (tree[0] === 'func_def') {
    const block = generateFuncDef(tree, 1, gctx);
    return block;
  }

  // ... 省略 ...
}

function generateFuncDef(tree, indent, gctx) {
  // tree = ["func_def", name, args[], body[]]
  // -- append to global context --
  // function hash: funcList['func1'] = [func_type, func_symbol, ret_type, func_block]
  //  ex) funcList['add'] = ['user_defined', '$add', 'i32', '...']

  const funcName = tree[1];
  const args = tree[2];
  const argCount = getLength(tree[2]);
  const funcSymbol = '$' + funcName;
  const funcType = 'i32';
  const funcResult = '(result ' + funcType + ')';

  let funcBlock = TABs(indent) + '(func ' + funcSymbol;

  // -- add temporary with empty body ---
  addGlobalFunc(gctx, funcName, funcSymbol, funcType, '');

  // --- agrs ---
  let argBlock = '';
  let i = 0;
  while (i < argCount) {
    argBlock = argBlock + ' ' + '(param $' + args[i] + ' i32)';
    i = i + 1;
  }
  funcBlock = funcBlock + argBlock + ' ' + funcResult + LF();

  // --- func bocy ---
  // --- prepare new local context for inside of function --
  let newLctx = initialLocalContext();
  // --- NEED to add args as localVar
  i = 0;
  let name;
  let varName;
  while (i < argCount) {
    name = args[i];
    varName = '$' + name;
    newLctx[name] = varName;
    i = i + 1;
  }
  let funcBody = generate(tree[3], indent + 1, gctx, newLctx);

  // --- build func block ---
  const varOffset = argCount;
  const varBlock = generateVariableBlock(indent + 1, newLctx, varOffset);
  funcBlock = funcBlock + varBlock + funcBody;

  // --- add dummy return ---
  funcBlock = funcBlock + TABs(indent + 1) + 'i32.const 88' + LF();
  funcBlock = funcBlock + TABs(indent + 1) + 'return' + LF();

  // --- close func definition ---
  funcBlock = funcBlock + TABs(indent) + ')' + LF();

  // ==== whole func definition ===
  addGlobalFunc(gctx, funcName, funcSymbol, funcType, funcBlock);
  println('---funcBlock--');
  println(funcBlock);

  // --- no code in this timing --
  const empty = '';
  return empty;
}

関数定義を生成する generateFuncDef() では、複数の処理を行います。

• 前回用意したグローバルコンテキストに関数の情報を仮登録する ... addGlobalFunc()
• 関数のシンボル、引数、戻り値を宣言を生成
• 関数内で使うローカル変数を宣言生成　... generateVariableBlock()
• 関数の処理の中身を生成
  • この時に、関数内のローカル変数を扱うために新しいローカルコンテキストを生成
  • 手抜きで関数の最後で return しないケース（if-then, while のブロック内で return)をきちんと判別できていないため、ダミーの return 文を最後に付与
• 改めて、グローバルコンテキストに関数の情報を登録する ... addGlobalFunc()

最初に関数の情報を仮登録するのは、以前のLLVMを使ったコンパイラーを実装したときに、ユーザー定義関数の中で再帰呼び出しを行うケースに対処した経験に基づいています。

function generateVariableBlock(indent, lctx, offset) {
  const variables = getKeys(lctx);
  const len = getLength(variables);
  let key;
  let i = offset;
  let block = '';
  let varName;
  while (i < len) {
    key = variables[i];
    varName = lctx[key];

    //  (local $a i32)
    block = block + TABs(indent) + '(local ' + varName + ' i32)' + LF();

    i = i + 1;
  }

  return block;
}

//ex) funcList['add'] = ['user_defined', '$add', 'i32', '.....']
function addGlobalFunc(gctx, name, symbol, type, funcBlock) {
  let funcList = gctx['funcList'];
  funcList[name] = ['user_defined', symbol, type, funcBlock];
}

function getGlobalFunc(gctx, name) {
  const funcList = gctx['funcList'];
  return funcList[name];
}

また今回扱うユーザ定義に関数には、いくつか制約があります。

• 呼び出すよりも前に定義されている必要がある
• 必ず 32ビット符号付整数を返す

関数からのreturn

ユーザ定義関数の生成のために、return文にも対応しておきます。return文は次の単純化ASTで表現されます。

[ 'ret', _戻り値の式_ ]

これに対応する様に、コンパイラーを拡張します。

function generate(tree, indent, gctx, lctx) {
  // ... 省略 ...

  // --- return from function ---
  if (tree[0] === 'ret') {
    const block = generateFuncRet(tree, indent, gctx, lctx);
    return block;
  }

  // ... 省略 ...
}

function generateFuncRet(tree, indent, gctx, lctx) {
  const valueBlock = generate(tree[1], indent, gctx, lctx);
  let block = valueBlock + LF();
  block = block + TABs(indent) + 'return' + LF();

  return block;
}

generateFuncRet()では、戻り値の式の結果をスタックに積み、return を行います。

ユーザ定義関数の呼び出し

関数呼び出しの部分は、もともと用意していた組み込み関数のputn()/puts()呼び出しの部分を拡張しました。

function generate(tree, indent, gctx, lctx) {
  // ... 省略 ...

  if (tree[0] === 'func_call') {  // tree = ['func_call', 'name', arg1, arg2, ... ]
    const funcName = tree[1];

    // --- builtin function ---
    if (funcName === 'putn') {
      return generateCallPutn(tree, indent, gctx, lctx);
    }
    if (funcName === 'puts') {
      return generateCallPuts(tree, indent, gctx, lctx);
    }

    // --- user defined functions ---
    const block = generateFunCall(tree, indent, gctx, lctx);
    return block;
  }

  // ... 省略 ...
}

// --- user defined functions ---
function generateFunCall(tree, indent, gctx, lctx) {
  // tree ['func_call', 'name', arg1, arg2, ... ]
  const funcName = tree[1];
  const gfunc = getGlobalFunc(gctx, funcName);
  // gfunc : ['user_defined', symbol, type, funcBlock];
  if (gfunc) {
    let block = '';

    // --- args ---
    let argBlock = '';
    let i = 0;
    while (tree[2 + i]) {
      argBlock = argBlock + generate(tree[2 + i], indent, gctx, lctx) + LF();
      i = i + 1;
    }

    // --- call ---
    block = block + argBlock;
    block = block + TABs(indent) + 'call ' + gfunc[1] + LF();

    return block;
  }

  println('-- ERROR: unknown func in generateFunCall() name=' + funcName);
  abort();
}

実際の呼び出し処理は generateFunCall() で生成しています。

• getGlobalFunc() で、登録済みのユーザ定義関数を探し、シンボル名を取得
• 引数をスタックに積む
• 関数を呼び出す

定義時の引数の数と、呼び出し時に渡される引数の数のチェックはしていません。

ユーザ定義関数の実体の生成

関数定義の処理では、グローバルコンテキストに登録するとこまでしかやっていないので、最後に全体を生成するさに、一緒に出力します。

// ---- compile simplified tree into WAT ---
function compile(tree, gctx, lctx) {
  // ... 省略 ...

  // --- export main function  --
  block = block + LF();
  block = block + TAB() + ';; ---- export main function  ---' + LF();
  block = block + TAB() + '(export "exported_main" (func $main))' + LF();
  block = block + TAB() + '(func $main (result i32)' + LF();
  block = block + varBlock + LF();
  block = block + mainBlock + LF();
  block = block + TAB() + TAB() + 'return' + LF();
  block = block + TAB() + ')' + LF();

  // ---- global user_defined functions ---
  block = block + generateGlobalFunctions(gctx);

  block = block + ')';

  return block;
}

function generateGlobalFunctions(gctx) {
  let block = LF();
  block = block + TAB() + ';; --- user_defined functions ---' + LF();

  const names = getGlobalFunctionNames(gctx);
  const len = getLength(names);
  println('--getGlobalFunctionNames() len=' + len);
  let key;
  let i = 0;
  let gfunc;
  while (i < len) {
    key = names[i];
    gfunc = getGlobalFunc(gctx, key);
    // gfunc : ['user_defined', symbol, type, funcBlock];
    if (gfunc[0] === 'user_defined') {
      block = block + gfunc[3] + LF();
    }

    i = i + 1;
  }

  return block;
}

generateGlobalFunctions() では、登録されているユーザ定義関数の実体を取り出して、順番に出力しています。

ユーザ定義関数の利用例：再帰でフィボナッチ数列

今回の目標の一つだったフィボナッチ数列の計算を、再帰呼び出しを使ってやってみます。

元のソース

fib_func.js

function fib(x) {
  if (x <= 1) {
    return x
  }
  else {
    return fib(x - 1) + fib(x - 2);
  }
}

let i = 0;
while (i < 10) {
  putn(fib(i));
  i = i + 1;
}

0;

コンパイラーで生成された WAT

generated.wat

(module
  ;; ---- builtin func imports ---
  (func $putn (import "imports" "imported_putn") (param i32))
  (func $puts (import "imports" "imported_puts") (param i32))

  ;; ---- export static string  ---
  (memory $string_area 1) ;; string_area 64KiB
  (export "exported_string" (memory $string_area))

  ;; ---- export main function  ---
  (export "exported_main" (func $main))
  (func $main (result i32)
    (local $i i32)

    i32.const 0
    set_local $i

    loop ;; --begin of while loop--
      get_local $i
      i32.const 10
      i32.lt_s
      if
        get_local $i
        call $fib

        call $putn

        get_local $i
        i32.const 1
        i32.add

        set_local $i

        br 1 ;; --jump to head of while loop--
      end ;; end of if-then
    end ;; --end of while loop--
    i32.const 0

    return
  )

  ;; --- user_defined functions ---
  (func $fib (param $x i32) (result i32)
    get_local $x
    i32.const 1
    i32.le_s

    if
      get_local $x
      return
    else
      get_local $x
      i32.const 1
      i32.sub

      call $fib

      get_local $x
      i32.const 2
      i32.sub

      call $fib

      i32.add

      return
    end
    i32.const 88
    return
  )

)

実行結果

$ wat2wasm generated.wat
$ node run_wasm_builtin.js generated.wasm
Loading wasm file: generated.wasm
0
1
1
2
3
5
8
13
21
34
ret code=0

無事ユーザ定義関数を使い、実行することができました。

次回は

WASMコンパイラーの実装は今回で一区切りなのですが、追加の取り組みとして WASI (WebAssembly System Interface) 対応を行う予定です。

ここまでのソース

GitHubにソースを上げておきます。

• GitHubのレポジトリ ... https://github.com/mganeko/mini_node_wasm
• mininode_wasm_07.js ... 今回のWASMコンパイラー
• mininode_wasm_08.js ... 戻り値でエラーが出るケースに対処した最終的なWASMコンパイラー
  • ※内容の説明は省略します（よく思い出せない）
• run_wasm_builtin.js ... WASMの実行につかうファイル
  sample/fizzbuzz_func.js ... ユーザ定義関数を利用したFizzBuzzのサンプル
  sample/fib_func.js ... ユーザ定義関数を利用したフィボナッチ数列のサンプル