#!/usr/bin/env python2

'''
Fast static linker for emscripten outputs. Specifically this links asm.js modules.

Usage: emlink.py [main module] [side module] [output name]

  Main module should be built with -s MAIN_MODULE=1 (and possibly -s DISABLE_GL_EMULATION=1, if you do not need GL emulation)
  Side module should be built with -s SIDE_MODULE=1

Note that the output file can be used as a main module, so you can link multiple
side modules into a main module that way.

Limitations:

 * Modules cannot be minified (but can be minified after linking)
 * We duplicate code in some cases, like overlapping names in different modules, and function aliases
 * We do not handle sharing of global constants across modules, only code (you can make a function to
   access a constant, if you need that)

Overall, this linking approach should be fast to perform, but generate less-optimal results than
to link all the bitcode together and build to JS as a single project. Final builds should be
done in that approach, but during development static linking can avoid recompiling your entire
project all the time.
'''

import os, subprocess, sys, re
from tools import shared
from tools import js_optimizer

try:
  me, main, side, out = sys.argv[:4]
except:
  print >> sys.stderr, 'usage: emlink.py [main module] [side module] [output name]'
  sys.exit(1)

print 'Main module:', main
print 'Side module:', side
print 'Output:', out

class AsmModule():
  def __init__(self, filename):
    self.filename = filename
    self.js = open(filename).read()

    self.start_asm = self.js.find(js_optimizer.start_asm_marker)
    self.start_funcs = self.js.find(js_optimizer.start_funcs_marker)
    self.end_funcs = self.js.rfind(js_optimizer.end_funcs_marker)
    self.end_asm = self.js.rfind(js_optimizer.end_asm_marker)

    # pre
    self.pre_js = self.js[:self.start_asm]

    # heap initializer
    mem_init = re.search(shared.JS.memory_initializer_pattern, self.pre_js)
    if mem_init:
      self.mem_init_full_js = mem_init.group(0)
      self.mem_init_js = mem_init.groups(0)[0][:-2]
      self.mem_init_size = self.mem_init_js.count(',') + self.mem_init_js.count('concat') + 1 # XXX add testing for large and small ones
      pad = 8 - (self.mem_init_size % 8)
      #print >> sys.stderr, 'pad', self.mem_init_size, pad
      if pad < 8:
        self.mem_init_js += '.concat([' + ','.join(['0']*pad) + '])'
        self.mem_init_size += pad
    else:
      self.mem_init_js = ''
      self.mem_init_size = 0
    #print >> sys.stderr, self.mem_init_js

    # global initializers
    global_inits = re.search(shared.JS.global_initializers_pattern, self.pre_js)
    if global_inits:
      self.global_inits_js = global_inits.group(0)
      self.global_inits = map(lambda init: init.split('{')[2][1:].split('(')[0], global_inits.groups(0)[0].split(','))
    else:
      self.global_inits_js = ''
      self.global_inits = []

    # imports
    self.imports_js = self.js[self.start_asm:self.start_funcs]
    self.imports = [m.group(0) for m in js_optimizer.import_sig.finditer(self.imports_js)]
    #print 'imports', self.imports

    # funcs
    self.funcs_js = self.js[self.start_funcs:self.end_funcs]
    self.funcs = [m.group(2) for m in js_optimizer.func_sig.finditer(self.funcs_js)]
    #print 'funcs', self.funcs

    # tables and exports
    post_js = self.js[self.end_funcs:self.end_asm]
    ret = post_js.find('return')
    self.tables_js = post_js[:ret]
    self.exports_js = post_js[ret:]
    self.tables = self.parse_tables(self.tables_js)
    self.exports = set([export.strip() for export in self.exports_js[self.exports_js.find('{')+1:self.exports_js.find('}')].split(',')])

    # post
    self.post_js = self.js[self.end_asm:]
    self.module_defs = set(re.findall('var [\w\d_$]+ = Module\["[\w\d_$]+"\] = asm\["[\w\d_$]+"\];\n', self.post_js))

  def relocate_into(self, main):
    # heap initializer
    concat = '.concat(' if main.mem_init_js and self.mem_init_js else ''
    end = ')' if main.mem_init_js and self.mem_init_js else ''
    allocation = main.mem_init_js + concat + self.mem_init_js + end
    if allocation:
      full_allocation = '/* memory initializer */ allocate(' + allocation + ', "i8", ALLOC_NONE, Runtime.GLOBAL_BASE)'
      main.pre_js = re.sub(shared.JS.memory_initializer_pattern if main.mem_init_js else shared.JS.no_memory_initializer_pattern, full_allocation, main.pre_js, count=1)
      main.pre_js = re.sub('STATICTOP = STATIC_BASE \+ (\d+);', 'STATICTOP = STATIC_BASE + %d' % (main.mem_init_size + side.mem_init_size), main.pre_js, count=1)

    # imports
    main_imports = set(main.imports)
    new_imports = [imp for imp in self.imports if imp not in main_imports]
    main.imports_js += '\n'.join(new_imports)

    # Find function name replacements TODO: do not rename duplicate names with duplicate contents, just merge them
    main_funcs = set(main.funcs)
    replacements = {}
    for i in range(len(self.funcs)):
      rep = func = self.funcs[i]
      while rep in main_funcs:
        rep += '_'
        replacements[func] = rep
    #print >> sys.stderr, 'replacements:', replacements

    # tables
    f_bases = {}
    f_sizes = {}
    for table, data in self.tables.iteritems():
      main.tables[table] = self.merge_tables(table, main.tables.get(table), data, replacements, f_bases, f_sizes)
    main.combine_tables()

    # relocate
    temp = shared.Building.js_optimizer(self.filename, ['asm', 'relocate'], extra_info={
      'replacements': replacements,
      'fBases': f_bases,
      'hBase': main.mem_init_size
    })
    #print >> sys.stderr, 'relocated side into', temp
    relocated_funcs = AsmModule(temp)
    shared.try_delete(temp)
    main.extra_funcs_js = relocated_funcs.funcs_js.replace(js_optimizer.start_funcs_marker, '\n')

    # update function table uses
    def update_fts(what):
      def fix(m):
        table = 'FUNCTION_TABLE_' + m.group(1)
        contents = m.group(2)
        assert '[' not in contents # TODO handle nesting
        return '%s[%s&%d]' % (table, contents, f_sizes[table]-1)
      return re.sub('FUNCTION_TABLE_(\w+)\[(.*)& ?(\d+)\]', fix, what)
    main.funcs_js = update_fts(main.funcs_js)
    main.extra_funcs_js = update_fts(main.extra_funcs_js)

    # global initializers
    if self.global_inits:
      my_global_inits = map(lambda init: replacements[init] if init in replacements else init, self.global_inits)
      all_global_inits = map(lambda init: '{ func: function() { %s() } }' % init, main.global_inits + my_global_inits)
      all_global_inits_js = '/* global initializers */ __ATINIT__.push(' + ','.join(all_global_inits) + ');'
      if main.global_inits:
        target = main.global_inits_js
      else:
        target = '// === Body ===\n'
        all_global_inits_js = target + all_global_inits_js
      main.pre_js = main.pre_js.replace(target, all_global_inits_js)

    # exports
    def rep_exp(export):
      key, value = export.split(':')
      if key in replacements:
        repped = replacements[key]
        return repped + ': ' + repped
      return export
    my_exports = map(rep_exp, self.exports)
    exports = main.exports.union(my_exports)
    main.exports_js = 'return {' + ','.join(list(exports)) + '};\n})\n'

    # post
    def rep_def(deff):
      key = deff.split(' ')[1]
      if key in replacements:
        rep = replacements[key]
        return 'var %s = Module["%s"] = asm["%s"];\n' % (rep, rep, rep)
      return deff
    my_module_defs = map(rep_def, self.module_defs)
    new_module_defs = set(my_module_defs).difference(main.module_defs)
    if len(new_module_defs) > 0:
      position = main.post_js.find('Runtime.') # Runtime is the start of the hardcoded ones
      main.post_js = main.post_js[:position] + ''.join(list(new_module_defs)) + '\n' + main.post_js[position:]

  def write(self, out):
    f = open(out, 'w')
    f.write(self.pre_js)
    f.write(self.imports_js)
    f.write(self.funcs_js)
    f.write(self.extra_funcs_js)
    f.write(self.tables_js)
    f.write(self.exports_js)
    f.write(self.post_js)
    f.close()

  # Utilities

  def parse_tables(self, js):
    tables = {}
    parts = js.split(';')
    for part in parts:
      if '=' not in part: continue
      part = part.split('var ')[1]
      name, data = part.split(' = ')
      tables[name] = data
    return tables

  def merge_tables(self, table, main, side, replacements, f_bases, f_sizes):
    side = side[1:-1].split(',')
    side = map(lambda f: replacements[f] if f in replacements else f, side)
    if not main:
      f_bases[table] = 0
      f_sizes[table] = len(side)
      return '[' + ','.join(side) + ']'
    main = main[1:-1].split(',')
    # TODO: handle non-aliasing case too
    assert len(main) % 2 == 0
    f_bases[table] = len(main)
    ret = main + side
    size = 2
    while size < len(ret): size *= 2
    ret = ret + [0]*(size - len(ret))
    assert len(ret) == size
    f_sizes[table] = size
    return '[' + ','.join(ret) + ']'

  def combine_tables(self):
    self.tables_js = '// EMSCRIPTEN_END_FUNCS\n'
    for table, data in self.tables.iteritems():
      self.tables_js += 'var %s = %s;\n' % (table, data)

main = AsmModule(main)
side = AsmModule(side)

side.relocate_into(main)
main.write(out)