text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. (http://www.facebook.com)
  
  # Add library directory so we can import the compiler.
  import os
  import sys
  
  SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
  sys.path.append(os.path.join(SCRIPT_DIR, "..", "library"))
  
  import argparse
  import marshal
  import re
  import tempfile
  from collections import namedtuple
  from importlib.util import MAGIC_NUMBER
  from pathlib import Path
  from types import CodeType
  
  from _compiler import PyroCodeGenerator, compile as compiler_compile
  
  
  # While we do use a project specific compiler here, we still use the marshal
  # module of the host runtime and need to make sure it matches our expectations.
  assert (
      MAGIC_NUMBER == (3413).to_bytes(2, "little") + b"\r\n"
  ), "unexpected marshalling format"
  
  
  INIT_MODULE_NAME = "__init_module__"
  
  
  # Pyro-specific flags to mark code objects (see also `runtime/objects.h`)
  CO_BUILTIN = 0x2000000
  CO_METADATA = 0x4000000
  
  
  def builtin_function_template():
      _builtin()  # noqa: F821
  
  
  def builtin_function_template_with_docu():
      """template with docu comment"""
      _builtin()  # noqa: F821
  
  
  def marked_with_intrinsic(code):
      consts = code.co_consts
      return (
          consts
          and isinstance(consts, tuple)
          and isinstance(consts[0], str)
          and consts[0].startswith("$intrinsic$")
      )
  
  
  def builtin_cpp_id(module_name, fqname, suffix=""):
      if len(fqname) == 1:
          return f"FUNC({module_name}, {fqname[0]}{suffix})"
      if len(fqname) == 2:
          return f"METH({fqname[0]}, {fqname[1]}{suffix})"
      raise Exception(f"Don't know how to create C++ name for {fqname}")
  
  
  def builtin_intrinsic_id(module_name, fqname):
      if len(fqname) == 1:
          return f"FUNC({module_name}, {fqname[0]}_intrinsic)"
      if len(fqname) == 2:
          return f"METH({fqname[0]}, {fqname[1]}_intrinsic)"
      raise Exception(f"Don't know how to create C++ name for {fqname}")
  
  
  def mark_native_functions(code, builtins, intrinsics, module_name, fqname=()):
      """Checks whether the code object or recursively any code objects in the
      constants tuple only contains a `_builtin()` call and should be marked
      as a native function. If so an entry is appended to the `builtins` list
      and the code object is marked with the CO_BUILTIN flag."""
  
      flags = code.co_flags
      metadata = [0]
      if marked_with_intrinsic(code):
          flags |= CO_METADATA
          intrinsic_id = builtin_intrinsic_id(module_name, fqname)
          intrinsics.append(intrinsic_id)
          # The intrinsic IDs are biased by 1 so that 0 means no intrinsic
          metadata[0] = len(intrinsics)
  
      if code.co_names == ("_builtin",) and (
          code.co_code == builtin_function_template.__code__.co_code
          or code.co_code == builtin_function_template_with_docu.__code__.co_code
      ):
          assert code.co_stacksize == 1
          assert code.co_freevars == ()
          assert code.co_cellvars == ()
          # INIT_MODULE_NAME is magic and should not be used.
          assert code.co_name != INIT_MODULE_NAME
  
          builtin_id = builtin_cpp_id(module_name, fqname)
          builtin_num = len(builtins)
          builtins.append(builtin_id)
  
          # We abuse the stacksize field to store the index of the builtin
          # function. This way we do not need to invent a new marshaling type
          # just for this.
          flags |= CO_BUILTIN
          consts = (tuple(metadata),) if flags & CO_METADATA else ()
          new_code = code.replace(
              co_stacksize=builtin_num,
              co_flags=flags,
              co_code=b"",
              co_consts=consts,
              co_names=(),
              co_filename="",
              co_firstlineno=0,
              co_lnotab=b"",
              co_freevars=(),
              co_cellvars=(),
          )
          return new_code
  
      # Recursively transform code objects in the constants tuple.
      changed = flags & CO_METADATA
      new_consts = [tuple(metadata)] if changed else []
      for c in code.co_consts:
          if isinstance(c, CodeType):
              assert c.co_name != "<module>"
              sub_name = fqname + (c.co_name,)
              new_c = mark_native_functions(
                  c, builtins, intrinsics, module_name, sub_name
              )
              changed |= new_c is not c
          else:
              new_c = c
          new_consts.append(new_c)
      if not changed:
          return code
  
      new_consts = tuple(new_consts)
      return code.replace(co_flags=flags, co_consts=new_consts)
  
  
  def to_char_array(byte_array):
      """Returns a string with a brace-enclosed list of C char array elements"""
      members = ", ".join(f"0x{int(byte):02x}" for byte in byte_array)
      return f"{{{members}}}"
  
  
  ModuleData = namedtuple(
      "ModuleData", ("name", "initializer", "builtin_init", "is_package")
  )
  
  
  def get_module_name(filename):
      path = Path(filename)
      module_name = path.name
      assert module_name.endswith(".py")
      module_name = module_name[:-3]
      for parent in path.parents:
          if not os.path.exists(Path(parent, "__init__.py")):
              break
          module_name = f"{parent.name}.{module_name}"
      return module_name
  
  
  def process_module(filename, builtins, intrinsics):
      fullname = get_module_name(filename)
      if fullname.endswith(".__init__"):
          fullname = fullname[:-9]
          is_package = True
      else:
          is_package = False
      with open(filename) as fp:
          source = fp.read()
      builtin_init = "$builtin-init-module$" in source
      flags = PyroCodeGenerator.consts.CO_FUTURE_ANNOTATIONS
      module_code = compiler_compile(
          source,
          filename,
          "exec",
          flags,
          dont_inherit=None,
          optimize=0,
      )
      marked_code = mark_native_functions(module_code, builtins, intrinsics, fullname)
      # We don't write pyc headers because it would make bootstrapping tricky.
      bytecode = marshal.dumps(marked_code)
      initializer = to_char_array(bytecode)
      return ModuleData(fullname, initializer, builtin_init, is_package)
  
  
  def gen_frozen_modules_cpp(modules):
      result = f"""\
  #include "frozen-modules.h"
  
  #include "builtins.h"
  #include "globals.h"
  #include "modules.h"
  
  namespace py {{
  """
  
      for module in modules:
          module_ident = module.name.replace(".", "_")
          result += f"""
  static const byte kData_{module_ident}[] = {module.initializer};
  """
  
      result += f"""
  const FrozenModule kFrozenModules[] = {{
  """
  
      for module in modules:
          module_ident = module.name.replace(".", "_")
          is_package = "true" if module.is_package else "false"
          init = (
              f"FUNC({module.name}, {INIT_MODULE_NAME})"
              if module.builtin_init
              else "nullptr"
          )
          result += f"""\
      {{"{module.name}", ARRAYSIZE(kData_{module_ident}), kData_{module_ident}, {init}, {is_package}}},
  """
  
      result += f"""\
  }};
  const word kNumFrozenModules = ARRAYSIZE(kFrozenModules);
  
  }}  // namespace py
  """
      return result
  
  
  def gen_frozen_modules_h():
      return f"""\
  #include "modules.h"
  
  namespace py {{
  
  extern const FrozenModule kFrozenModules[];
  extern const word kNumFrozenModules;
  
  }}  // namespace py
  """
  
  
  def gen_builtins_cpp(builtins, intrinsics):
      result = f"""\
  #include "builtins.h"
  
  namespace py {{
  
  const BuiltinFunction kBuiltinFunctions[] = {{
  """
  
      for builtin in builtins:
          result += f"""\
      {builtin},
  """
  
      result += """\
  };
  
  const IntrinsicFunction kIntrinsicFunctions[] = {
  """
      for intrinsic in intrinsics:
          result += f"""\
      {intrinsic},
  """
  
      result += """\
  };
  const word kNumBuiltinFunctions = ARRAYSIZE(kBuiltinFunctions);
  const word kNumIntrinsicFunctions = ARRAYSIZE(kIntrinsicFunctions);
  }  // namespace py
  """
      return result
  
  
  def gen_builtins_h(modules, builtins, intrinsics):
      result = f"""\
  #include "globals.h"
  #include "handles-decl.h"
  #include "modules.h"
  #include "view.h"
  
  namespace py {{
  
  class Thread;
  
  extern const BuiltinFunction kBuiltinFunctions[];
  extern const IntrinsicFunction kIntrinsicFunctions[];
  extern const word kNumBuiltinFunctions;
  extern const word kNumIntrinsicFunctions;
  
  """
  
      for builtin in builtins:
          result += f"""\
  RawObject {builtin}(Thread* thread, Arguments args) ALIGN_16;
  """
      result += "\n"
      for intrinsic in intrinsics:
          result += f"""\
  bool {intrinsic}(Thread* thread) ALIGN_16;
  """
      result += "\n"
      for module in modules:
          if module.builtin_init:
              result += f"""\
  void FUNC({module.name}, {INIT_MODULE_NAME})(Thread*, const Module&, View<byte>);
  """
  
      result += f"""
  }}  // namespace py
  """
      return result
  
  
  def write_if_different(path, string):
      try:
          if path.read_text() == string:
              return
      except OSError:
          pass
      with tempfile.NamedTemporaryFile(dir=path.parent, delete=False, mode="w+t") as fp:
          fp.write(string)
      Path(fp.name).replace(path)
  
  
  def main(argv):
      parser = argparse.ArgumentParser(prog=argv[0])
      parser.add_argument("dir", help="directory where output is written")
      parser.add_argument("files", nargs="+", help="list of files to freeze")
      args = parser.parse_args(argv[1:])
      os.makedirs(args.dir, exist_ok=True)
      builtins = []
      intrinsics = []
      modules = [
          process_module(filename, builtins, intrinsics) for filename in args.files
      ]
      modules.sort(key=lambda x: x.name)
      write_if_different(
          Path(args.dir, "builtins.cpp"), gen_builtins_cpp(builtins, intrinsics)
      )
      write_if_different(
          Path(args.dir, "builtins.h"), gen_builtins_h(modules, builtins, intrinsics)
      )
      write_if_different(
          Path(args.dir, "frozen-modules.cpp"), gen_frozen_modules_cpp(modules)
      )
      write_if_different(Path(args.dir, "frozen-modules.h"), gen_frozen_modules_h())
  
  
  if __name__ == "__main__":
      main(sys.argv)
  