<commit_before>from atom import *
from base import *
import compiler
from compiler import ast

ScopeContext = DT('ScopeContext', ('indent', int),
                                  ('syms', {(str, bool): object}),
                                  ('prevContext', None))
SCOPE = new_env('SCOPE', ScopeContext)

OmniContext = DT('OmniContext', ('imports', [object]),
                                ('exports', [object]),
                                ('missingRefs', {(str, bool): [object]}),
                                ('loadedDeps', set([Module])),
                                ('directlyImportedModuleNames', set([str])))
OMNI = new_env('OMNI', OmniContext)

AstType = new_extrinsic('AstType', Type)
AstHint = new_extrinsic('AstHint', {str: str})

def ast_envs():
    omni = OmniContext({}, {}, {}, set(), set())
    scope = ScopeContext(0, {}, None)
    return omni, scope

def is_top_level():
    return env(SCOPE).indent == 0

loaded_module_export_names = {}

valueNamespace = 'value'
typeNamespace = 'type'
symbolNamespace = 'symbol'

def identifier(obj, name=None, namespace=valueNamespace,
               permissible_nms=frozenset(), export=False):
    scope = env(SCOPE)
    already_Named = name is None
    if already_Named:
        name = extrinsic(Name, obj)
    key = (name, namespace)
    assert name in permissible_nms or key not in scope.syms, (
            "Symbol '%s' already in %s env\n%s" % (
            name, namespace,
            '\n'.join('\t'+str(s) for s, ns in scope.syms if ns == namespace)))
    try:
        b = bind_kind(obj)
    except ValueError:
        assert False, "Can't bind %s %r to %r" % (namespace, obj, name)
    scope.syms[key] = (obj, b)
    if not already_Named:
        add_extrinsic(Name, obj, name)

    omni = env(OMNI)
    missing_binds = omni.missingRefs.get(key)
    if missing_binds is not None:
        for bind in missing_binds:
            bind.binding = b(obj)
        del omni.missingRefs[key]
    if export:
        omni.exports[key] = (obj, b)

def ident_exists(nm, namespace=valueNamespace):
    scope = env(SCOPE)
    key = (nm, namespace)
    while scope is not None:
        o = scope.syms.get(key)
        if o is not None:
            var, b = o
            return Just(b(var))
        scope = scope.prevContext

    ext = env(OMNI).imports.get(key)
    if ext is not None:
        var, b = ext
        return Just(b(var))

    return Nothing()

def refs_existing(nm, namespace=valueNamespace):
    ref = ident_exists(nm, namespace)
    if isJust(ref):
        return Bind(fromJust(ref))
    key = (nm, namespace)
    fwd_ref = Bind(key)
    omni = env(OMNI)
    omni.missingRefs.setdefault(key, []).append(fwd_ref)
    return fwd_ref

def bind_kind(v):
    if isinstance(v, Var):
        return BindVar
    elif isinstance(v, Builtin):
        return BindBuiltin
    elif isinstance(v, Ctor):
        return BindCtor
    elif isinstance(v, (DataType, Env, Extrinsic)):
        return identity
    else:
        raise ValueError("Can't bind to %r" % (v,))

def type_ref(nm):
    if nm in types_by_name:
        return types_by_name[nm]()
    ref = ident_exists(nm, typeNamespace)
    return maybe_(TForward(nm), ref)

def destroy_forward_ref(ref):
    if not isinstance(ref.refAtom, basestring):
        return
    omni = env(OMNI)
    for t in (True, False):
        key = (ref.refAtom, t)
        refs = omni.missingRefs.get(key, [])
        if ref in refs:
            refs.remove(ref)
            if not refs:
                del omni.missingRefs[key]
            return
    assert False, "Couldn't find forward ref for destruction"

def inside_scope(f):
    def new_scope(*args, **kwargs):
        prev = env(SCOPE)
        new = ScopeContext(prev.indent + 1, {}, prev)
        return in_env(SCOPE, new, lambda: f(*args, **kwargs))
    new_scope.__name__ = f.__name__
    return new_scope

def unknown_stmt(node):
    print '??%s %s??' % (node.__class__,
          ', '.join(filter(lambda x: not x.startswith('_'), dir(node))))
    exit(-1)

def unknown_expr(node):
    print '%s?(%s)' % (str(node.__class__),
          ', '.join(filter(lambda x: not x.startswith('_'), dir(node))))

def make_grammar_decorator(default_dispatch):
    dispatch_index = {}
    def decorator(key):
        def inserter(f):
            dispatch_index[key] = f
            return f
        return inserter
    def dispatch(nd, *args):
        return dispatch_index.get(nd.__class__, default_dispatch)(nd, *args)
    return (decorator, dispatch)

def conv_type(t, tvars, dt=None):
    def unknown():
        assert isinstance(getattr(t, 'refAtom', None), basestring
                ), 'Unknown type: %r' % (t,)
        type_nm = t.refAtom
        destroy_forward_ref(t)
        return type_ref(type_nm)
    return match(t,
        ("BindBuiltin(_)", lambda: t),
        ("TPrim(_)", lambda: t),
        ("StrLit(s)", lambda s: parse_new_type(s, tvars)),
        ("ListLit([t])",
            lambda t: TArray(conv_type(t, tvars, dt))),
        ("_", unknown))

(top_level, conv_top_level) = make_grammar_decorator(unknown_stmt)
(stmt, conv_stmt) = make_grammar_decorator(unknown_stmt)
(expr, conv_expr) = make_grammar_decorator(unknown_expr)

def conv_exprs(elist):
    return map(conv_expr, elist)

# EXPRESSIONS

for (cls, op) in {ast.Add: '+', ast.Sub: '-',
                  ast.Mul: '*', ast.Mod: '%',
                  ast.Div: '/', ast.FloorDiv: '//',
                  ast.Bitand: '&', ast.Bitor: '|', ast.Bitxor: '^',
                  ast.LeftShift: '<<', ast.RightShift: '>>'}.iteritems():
    @expr(cls)
    def binop(e, o=op):
        return symcall(o, [conv_expr(e.left), conv_expr(e.right)])

for (cls, sym) in {ast.UnaryAdd: 'positive',
                   ast.UnarySub: 'negate',
                   ast.Not: 'not',
                   ast.Invert: 'invert'}.iteritems():
    @expr(cls)
    def unaop(e, s=sym):
        return symcall(s, [conv_expr(e.expr)])
del cls, sym

@expr(ast.And)
def conv_and(e):
    assert len(e.nodes) == 2
    return And(*map(conv_expr, e.nodes))

def unwrap_ast(node):
    if isinstance(node, ast.Const):
        return node.value
    elif isinstance(node, ast.Tuple):
        return tuple(map(unwrap_ast, node.nodes))
    else:
        assert False

def _make_dt(dt_nm, *args, **opts):
    dt_nm = unwrap_ast(dt_nm)
    if dt_nm not in DATATYPES:
        args = map(unwrap_ast, args)
        opts['maker'](dt_nm, *args)
    dt = DATATYPES[dt_nm].__form__
    identifier(dt, namespace=typeNamespace)
    for ctor in dt.ctors:
        identifier(ctor, export=True)
    return [TopDT(dt)]

def make_adt(*args):
    return _make_dt(*args, **(dict(maker=ADT)))
def make_dt(*args):
    return _make_dt(*args, **(dict(maker=DT)))

def make_env(nm, t):
    tvars = {}<commit_msg>Minor AST conversion fix<commit_after>from atom import *
from base import *
import compiler
from compiler import ast

ScopeContext = DT('ScopeContext', ('indent', int),
                                  ('syms', {(str, bool): object}),
                                  ('prevContext', None))
SCOPE = new_env('SCOPE', ScopeContext)

OmniContext = DT('OmniContext', ('imports', [object]),
                                ('exports', [object]),
                                ('missingRefs', {(str, bool): [object]}),
                                ('loadedDeps', set([Module])),
                                ('directlyImportedModuleNames', set([str])))
OMNI = new_env('OMNI', OmniContext)

AstType = new_extrinsic('AstType', Type)
AstHint = new_extrinsic('AstHint', {str: str})

def ast_envs():
    omni = OmniContext({}, {}, {}, set(), set())
    scope = ScopeContext(0, {}, None)
    return omni, scope

def is_top_level():
    return env(SCOPE).indent == 0

loaded_module_export_names = {}

valueNamespace = 'value'
typeNamespace = 'type'
symbolNamespace = 'symbol'

def identifier(obj, name=None, namespace=valueNamespace,
               permissible_nms=frozenset(), export=False):
    scope = env(SCOPE)
    already_Named = name is None
    if already_Named:
        name = extrinsic(Name, obj)
    key = (name, namespace)
    assert name in permissible_nms or key not in scope.syms, (
            "Symbol '%s' already in %s env\n%s" % (
            name, namespace,
            '\n'.join('\t'+str(s) for s, ns in scope.syms if ns == namespace)))
    try:
        b = bind_kind(obj)
    except ValueError:
        assert False, "Can't bind %s %r to %r" % (namespace, obj, name)
    scope.syms[key] = (obj, b)
    if not already_Named:
        add_extrinsic(Name, obj, name)

    omni = env(OMNI)
    missing_binds = omni.missingRefs.get(key)
    if missing_binds is not None:
        for bind in missing_binds:
            bind.binding = b(obj)
        del omni.missingRefs[key]
    if export:
        omni.exports[key] = (obj, b)

def ident_exists(nm, namespace=valueNamespace):
    scope = env(SCOPE)
    key = (nm, namespace)
    while scope is not None:
        o = scope.syms.get(key)
        if o is not None:
            var, b = o
            return Just(b(var))
        scope = scope.prevContext

    ext = env(OMNI).imports.get(key)
    if ext is not None:
        var, b = ext
        return Just(b(var))

    return Nothing()

def refs_existing(nm, namespace=valueNamespace):
    ref = ident_exists(nm, namespace)
    if isJust(ref):
        return Bind(fromJust(ref))
    key = (nm, namespace)
    fwd_ref = Bind(key)
    omni = env(OMNI)
    omni.missingRefs.setdefault(key, []).append(fwd_ref)
    return fwd_ref

def bind_kind(v):
    if isinstance(v, Var):
        return BindVar
    elif isinstance(v, Builtin):
        return BindBuiltin
    elif isinstance(v, Ctor):
        return BindCtor
    elif isinstance(v, (DataType, Env, Extrinsic)):
        return identity
    else:
        raise ValueError("Can't bind to %r" % (v,))

def type_ref(nm):
    if nm in types_by_name:
        return types_by_name[nm]()
    ref = ident_exists(nm, typeNamespace)
    return maybe_(TForward(nm), ref)

def destroy_forward_ref(ref):
    if not isinstance(ref.refAtom, basestring):
        return
    omni = env(OMNI)
    for t in (True, False):
        key = (ref.refAtom, t)
        refs = omni.missingRefs.get(key, [])
        if ref in refs:
            refs.remove(ref)
            if not refs:
                del omni.missingRefs[key]
            return
    assert False, "Couldn't find forward ref for destruction"

def inside_scope(f):
    def new_scope(*args, **kwargs):
        prev = env(SCOPE)
        new = ScopeContext(prev.indent + 1, {}, prev)
        return in_env(SCOPE, new, lambda: f(*args, **kwargs))
    new_scope.__name__ = f.__name__
    return new_scope

def unknown_stmt(node):
    print '??%s %s??' % (node.__class__,
          ', '.join(filter(lambda x: not x.startswith('_'), dir(node))))
    exit(-1)

def unknown_expr(node):
    print '%s?(%s)' % (str(node.__class__),
          ', '.join(filter(lambda x: not x.startswith('_'), dir(node))))

def make_grammar_decorator(default_dispatch):
    dispatch_index = {}
    def decorator(key):
        def inserter(f):
            dispatch_index[key] = f
            return f
        return inserter
    def dispatch(nd, *args):
        return dispatch_index.get(nd.__class__, default_dispatch)(nd, *args)
    return (decorator, dispatch)

def conv_type(t, tvars, dt=None):
    def unknown():
        assert isinstance(getattr(t, 'refAtom', None), basestring
                ), 'Unknown type: %r' % (t,)
        type_nm = t.refAtom
        destroy_forward_ref(t)
        return type_ref(type_nm)
    return match(t,
        ("BindBuiltin(_)", lambda: t),
        ("TPrim(_)", lambda: t),
        ("StrLit(s)", lambda s: parse_new_type(s, tvars)),
        ("ListLit([t])",
            lambda t: TArray(conv_type(t, tvars, dt))),
        ("_", unknown))

(top_level, conv_top_level) = make_grammar_decorator(unknown_stmt)
(stmt, conv_stmt) = make_grammar_decorator(unknown_stmt)
(expr, conv_expr) = make_grammar_decorator(unknown_expr)

def conv_exprs(elist):
    return map(conv_expr, elist)

# EXPRESSIONS

for (cls, op) in {ast.Add: '+', ast.Sub: '-',
                  ast.Mul: '*', ast.Mod: '%',
                  ast.Div: '/', ast.FloorDiv: '//',
                  ast.Bitand: '&', ast.Bitor: '|', ast.Bitxor: '^',
                  ast.LeftShift: '<<', ast.RightShift: '>>'}.iteritems():
    @expr(cls)
    def binop(e, o=op):
        return symcall(o, [conv_expr(e.left), conv_expr(e.right)])

for (cls, sym) in {ast.UnaryAdd: 'positive',
                   ast.UnarySub: 'negate',
                   ast.Not: 'not',
                   ast.Invert: 'invert'}.iteritems():
    @expr(cls)
    def unaop(e, s=sym):
        return symcall(s, [conv_expr(e.expr)])
del cls, sym

@expr(ast.And)
def conv_and(e):
    assert len(e.nodes) == 2
    return And(*map(conv_expr, e.nodes))

def unwrap_ast(node):
    if isinstance(node, ast.Const):
        return node.value
    elif isinstance(node, ast.Name):
        return node.name
    elif isinstance(node, ast.Tuple):
        return tuple(map(unwrap_ast, node.nodes))
    else:
        assert False

def _make_dt(dt_nm, *args, **opts):
    dt_nm = unwrap_ast(dt_nm)
    if dt_nm not in DATATYPES:
        args = map(unwrap_ast, args)
        opts['maker'](dt_nm, *args)
    dt = DATATYPES[dt_nm].__form__
    identifier(dt, namespace=typeNamespace)
    for ctor in dt.ctors:
        identifier(ctor, export=True)
    return [TopDT(dt)]

def make_adt(*args):
    return _make_dt(*args, **(dict(maker=ADT)))
def make_dt(*args):
    return _make_dt(*args, **(dict(maker=DT)))

def make_env(nm, t):
    tvars = {}
{'fraction_non_alphanumeric': 0.14125247779164526, 'fraction_numerical': 0.0007341604874825637, 'mean_word_length': 2.5063063063063065, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3887952', 'n_tokens_mistral': 4635, 'n_tokens_neox': 4529, 'n_words': 1098}