<commit_before>from . import compiler as C
from . import module as M
from . import types as T
from . import error as Q
from .ast import *
from collections import OrderedDict
from llvmlite import ir
import os.path


# Program structure ###########################################################

@program_node.method
def emit(self, module):
  module.exports = module.add_global(T.box, name=module.mangle('exports'))
  module.exports.initializer = static_table_alloc(module, name=module.mangle('exports.table'))

  imports = []
  links = []
  libs = []
  for stmt in self.stmts:
    ret = module.emit(stmt)
    if isinstance(stmt, import_node):
      imports.append(ret)
    elif isinstance(stmt, link_node):
      links.append(ret)
    elif isinstance(stmt, lib_node):
      libs.append(ret)

  return imports, links, libs


@program_node.method
def emit_main(self, module, mods=[]):
  with module.add_main():
    ret_ptr = module.alloc(T.box, T.null, name='ret_ptr')

    with module.add_abort() as abort:
      module.excall('rain_init_gc')
      module.excall('rain_init_args', *module.main.args)

      for tmp in mods:
        if 'init' in tmp:
          with tmp.borrow_builder(module):
            init_box = load_global(tmp, 'init')
            init_ptr = module.get_value(init_box, typ=T.vfunc(T.arg))
            module.check_callable(init_box, 0, unwind=module.catch)
            module.fncall(init_ptr, T.null, unwind=module.catch)

      module.excall('rain_main', ret_ptr, module['main'], unwind=module.catch)

      abort(module.builder.block)

      ret_code = module.excall('rain_box_to_exit', ret_ptr)
      module.builder.ret(ret_code)


@block_node.method
def emit(self, module):
  for stmt in self.stmts:
    module.emit(stmt)

  if self.expr:
    return module.emit(self.expr)

  return T.null


# Helpers #####################################################################

# Return the index to insert / fetch from a static table
# Note: if the key isn't found, the returned index points to a None constant
def static_table_idx(module, table_box, key_node):
  lpt_ptr = table_box.lpt_ptr
  arr_ptr = lpt_ptr.arr_ptr

  max = lpt_ptr.initializer.constant[1].constant
  items = arr_ptr.initializer.constant
  key_hash = key_node.hash()

  while True:
    if items[key_hash % max].constant is None:
      break

    if items[key_hash % max].key == key_node:
      break

    key_hash += 1

  return key_hash % max


# Insert a box into a static table
def static_table_put(module, table_box, key_node, val):
  key = module.emit(key_node)

  lpt_ptr = table_box.lpt_ptr

  if getattr(lpt_ptr, 'arr_ptr', None) is None:
    Q.abort('Global value is opaque.')

  arr_ptr = lpt_ptr.arr_ptr
  item = T.item([T.i32(1), key, val])
  item.key = key_node

  cur = lpt_ptr.initializer.constant[0].constant
  max = lpt_ptr.initializer.constant[1].constant
  items = arr_ptr.initializer.constant

  idx = static_table_idx(module, table_box, key_node)
  if items[idx].constant is None:
    cur += 1

  # TODO resize if cur > max / 2! currently, it infinite loops.

  items[idx] = item
  arr_ptr.initializer = arr_ptr.value_type(items)
  arr_gep = arr_ptr.gep([T.i32(0), T.i32(0)])

  lpt_ptr.initializer = lpt_ptr.value_type([T.i32(cur), T.i32(max), arr_gep])
  lpt_ptr.arr_ptr = arr_ptr


# Return a box from a static table
def static_table_get(module, table_box, key_node):
  lpt_ptr = table_box.lpt_ptr
  arr_ptr = lpt_ptr.arr_ptr
  items = arr_ptr.initializer.constant

  idx = static_table_idx(module, table_box, key_node)
  if items[idx].constant is None:
    return T.null

  return items[idx].constant[2]


# Allocate a static table
def static_table_alloc(module, name):
  arr_typ = T.arr(T.item, T.HASH_SIZE)
  arr_ptr = module.add_global(arr_typ, name=name + '.array')
  arr_ptr.initializer = arr_typ([None] * T.HASH_SIZE)
  arr_gep = arr_ptr.gep([T.i32(0), T.i32(0)])

  lpt_typ = T.lpt
  lpt_ptr = module.add_global(lpt_typ, name=name)
  lpt_ptr.initializer = lpt_typ([T.i32(0), T.i32(T.HASH_SIZE), arr_gep])
  lpt_ptr.arr_ptr = arr_ptr

  return static_table_from_ptr(module, lpt_ptr)


# Return a box from a static table
def static_table_from_ptr(module, ptr):
  box = T._table(ptr)
  box.lpt_ptr = ptr  # save this for later!
  return box


# Return a pointer to a static table's value box
def static_table_get_box_ptr(module, table_box, key_node):
  idx = static_table_idx(module, table_box, key_node)
  return table_box.lpt_ptr.arr_ptr.gep([T.i32(0), T.i32(idx), T.i32(2)])


# Repair a static table box from another one
def static_table_repair(new_box, old_box):
  if getattr(old_box, 'lpt_ptr', None):
    new_box.lpt_ptr = old_box.lpt_ptr


# Store a value into a global (respecting whether it's exported or not)
def store_global(module, name: str, value: "LLVM value"):
  if not isinstance(module[name], ir.GlobalVariable):
    table_box = module.exports.initializer
    key_node = str_node(name)
    static_table_put(module, table_box, key_node, value)

  module[name].initializer = value


# Load a value from a global (respecting whether it's exported or not)
def load_global(module, name: str):
  return module[name].initializer


# Simple statements ###########################################################

@assn_node.method
def emit(self, module):
  if isinstance(self.lhs, name_node):
    if module.is_global:
      if self.export:
        table_box = module.exports.initializer
        key_node = str_node(self.lhs.value)
        module[self.lhs.value] = static_table_get_box_ptr(module, table_box, key_node)

      if self.let:
        module[self.lhs] = module.find_global(T.box, name=module.mangle(self.lhs.value))
        module[self.lhs].linkage = '' # make sure we know it's visible here

      if self.lhs not in module:
        Q.abort("Undeclared global {!r}", self.lhs.value)

      store_global(module, self.lhs.value, module.emit(self.rhs))
      return

    if self.let:<commit_msg>emit: Add code gen for list unpacking<commit_after>from . import compiler as C
from . import module as M
from . import types as T
from . import error as Q
from .ast import *
from collections import OrderedDict
from llvmlite import ir
import os.path


# Program structure ###########################################################

@program_node.method
def emit(self, module):
  module.exports = module.add_global(T.box, name=module.mangle('exports'))
  module.exports.initializer = static_table_alloc(module, name=module.mangle('exports.table'))

  imports = []
  links = []
  libs = []
  for stmt in self.stmts:
    ret = module.emit(stmt)
    if isinstance(stmt, import_node):
      imports.append(ret)
    elif isinstance(stmt, link_node):
      links.append(ret)
    elif isinstance(stmt, lib_node):
      libs.append(ret)

  return imports, links, libs


@program_node.method
def emit_main(self, module, mods=[]):
  with module.add_main():
    ret_ptr = module.alloc(T.box, T.null, name='ret_ptr')

    with module.add_abort() as abort:
      module.excall('rain_init_gc')
      module.excall('rain_init_args', *module.main.args)

      for tmp in mods:
        if 'init' in tmp:
          with tmp.borrow_builder(module):
            init_box = load_global(tmp, 'init')
            init_ptr = module.get_value(init_box, typ=T.vfunc(T.arg))
            module.check_callable(init_box, 0, unwind=module.catch)
            module.fncall(init_ptr, T.null, unwind=module.catch)

      module.excall('rain_main', ret_ptr, module['main'], unwind=module.catch)

      abort(module.builder.block)

      ret_code = module.excall('rain_box_to_exit', ret_ptr)
      module.builder.ret(ret_code)


@block_node.method
def emit(self, module):
  for stmt in self.stmts:
    module.emit(stmt)

  if self.expr:
    return module.emit(self.expr)

  return T.null


# Helpers #####################################################################

# Return the index to insert / fetch from a static table
# Note: if the key isn't found, the returned index points to a None constant
def static_table_idx(module, table_box, key_node):
  lpt_ptr = table_box.lpt_ptr
  arr_ptr = lpt_ptr.arr_ptr

  max = lpt_ptr.initializer.constant[1].constant
  items = arr_ptr.initializer.constant
  key_hash = key_node.hash()

  while True:
    if items[key_hash % max].constant is None:
      break

    if items[key_hash % max].key == key_node:
      break

    key_hash += 1

  return key_hash % max


# Insert a box into a static table
def static_table_put(module, table_box, key_node, val):
  key = module.emit(key_node)

  lpt_ptr = table_box.lpt_ptr

  if getattr(lpt_ptr, 'arr_ptr', None) is None:
    Q.abort('Global value is opaque.')

  arr_ptr = lpt_ptr.arr_ptr
  item = T.item([T.i32(1), key, val])
  item.key = key_node

  cur = lpt_ptr.initializer.constant[0].constant
  max = lpt_ptr.initializer.constant[1].constant
  items = arr_ptr.initializer.constant

  idx = static_table_idx(module, table_box, key_node)
  if items[idx].constant is None:
    cur += 1

  # TODO resize if cur > max / 2! currently, it infinite loops.

  items[idx] = item
  arr_ptr.initializer = arr_ptr.value_type(items)
  arr_gep = arr_ptr.gep([T.i32(0), T.i32(0)])

  lpt_ptr.initializer = lpt_ptr.value_type([T.i32(cur), T.i32(max), arr_gep])
  lpt_ptr.arr_ptr = arr_ptr


# Return a box from a static table
def static_table_get(module, table_box, key_node):
  lpt_ptr = table_box.lpt_ptr
  arr_ptr = lpt_ptr.arr_ptr
  items = arr_ptr.initializer.constant

  idx = static_table_idx(module, table_box, key_node)
  if items[idx].constant is None:
    return T.null

  return items[idx].constant[2]


# Allocate a static table
def static_table_alloc(module, name):
  arr_typ = T.arr(T.item, T.HASH_SIZE)
  arr_ptr = module.add_global(arr_typ, name=name + '.array')
  arr_ptr.initializer = arr_typ([None] * T.HASH_SIZE)
  arr_gep = arr_ptr.gep([T.i32(0), T.i32(0)])

  lpt_typ = T.lpt
  lpt_ptr = module.add_global(lpt_typ, name=name)
  lpt_ptr.initializer = lpt_typ([T.i32(0), T.i32(T.HASH_SIZE), arr_gep])
  lpt_ptr.arr_ptr = arr_ptr

  return static_table_from_ptr(module, lpt_ptr)


# Return a box from a static table
def static_table_from_ptr(module, ptr):
  box = T._table(ptr)
  box.lpt_ptr = ptr  # save this for later!
  return box


# Return a pointer to a static table's value box
def static_table_get_box_ptr(module, table_box, key_node):
  idx = static_table_idx(module, table_box, key_node)
  return table_box.lpt_ptr.arr_ptr.gep([T.i32(0), T.i32(idx), T.i32(2)])


# Repair a static table box from another one
def static_table_repair(new_box, old_box):
  if getattr(old_box, 'lpt_ptr', None):
    new_box.lpt_ptr = old_box.lpt_ptr


# Store a value into a global (respecting whether it's exported or not)
def store_global(module, name: str, value: "LLVM value"):
  if not isinstance(module[name], ir.GlobalVariable):
    table_box = module.exports.initializer
    key_node = str_node(name)
    static_table_put(module, table_box, key_node, value)

  module[name].initializer = value


# Load a value from a global (respecting whether it's exported or not)
def load_global(module, name: str):
  return module[name].initializer


# Simple statements ###########################################################

@assn_node.method
def emit(self, module):
  if isinstance(self.lhs, list): # unpack
    if module.is_global:
      Q.abort("Unable to unpack at global scope")

    # evalute the RHS before storing anything
    src = module.emit(self.rhs)
    pieces = []
    for i, lhs in enumerate(self.lhs):
      piece = module.exfncall('rain_get', T.null, src, int_node(i).emit(module))
      pieces.append(module.load(piece))

    # store everything
    for lhs, piece in zip(self.lhs, pieces):
      if isinstance(lhs, name_node):
        if lhs not in module:
          Q.abort("Undeclared name {!r}", lhs.value)

        module[lhs].bound = True
        module.store(piece, module[lhs])

      elif isinstance(lhs, idx_node):
        table = module.emit(lhs.lhs)
        key = module.emit(lhs.rhs)
        module.exfncall('rain_put', table, key, piece)

      elif isinstance(lhs, list):
        pass # TODO

  elif isinstance(self.lhs, name_node):
    if module.is_global:
      if self.export:
        table_box = module.exports.initializer
        key_node = str_node(self.lhs.value)
        module[self.lhs.value] = static_table_get_box_ptr(module, table_box, key_node)

      if self.let:
        module[self.lhs] = module.find_global(T.box, name=module.mangle(self.lhs.value))
        module[self.lhs].linkage = '' # make sure we know it's visible here

      if self.lhs not in module:
        Q.abort("Undeclared global {!r}", self.lhs.value)

      store_global(module, self.lhs.value, module.emit(self.rhs))
      return

    if self.let:
{'fraction_non_alphanumeric': 0.14030412290327637, 'fraction_numerical': 0.006270575325286095, 'mean_word_length': 3.769719626168224, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2167320', 'n_tokens_mistral': 4587, 'n_tokens_neox': 4436, 'n_words': 1047}