#=========================================================================
# custom_exec.py
#=========================================================================
# The purpose of this custom_exec function is to turn dictionaries into
# ModuleDicts that results in much better JIT behavior under PyPy.
#
# Author : Shunning Jiang
# Date   : Feb 12, 2020

try:
  import __pypy__

  def _normalize_dict( _dict ):
    # Credit: Lin Cheng
    # return the original dictionary if not running on PyPy
    try:
      from __pypy__ import newdict, strategy
    except Exception:
      return _dict

    # return the original dictionary if already using ModuleDictStrategy
    if strategy( _dict ) == "ModuleDictStrategy":
      return _dict
    # create a new module dict
    new_dict = newdict("module")
    # copy over entries
    for key, value in _dict.items():
      new_dict[key] = value
    return new_dict

  def custom_exec( prog, _globals, _locals ):
    assert _globals is not None
    assert _locals  is not None

    norm_globals = _globals
    norm_locals  = _locals
    if _globals is not None:
      norm_globals = _normalize_dict( _globals )
    if _locals is not None:
      norm_locals  = _normalize_dict( _locals )

    exec( prog, norm_globals, norm_locals )
    # Local may have more stuff generated by the code.
    # We need to put this back to the original _locals
    if _locals is not None:
      _locals.update( norm_locals )

except Exception:
  custom_exec = exec
