<commit_before>from base import *
from bedrock import *
from globs import *
from hashlib import sha256
from os import system

SerialState = DT('SerialState',
        ('file', file),
        ('hash', None),
        ('index', int),
        ('deps', {Module: int}))

Serialize = new_env('Serialize', SerialState)

def _write(b):
    state = env(Serialize)
    state.file.write(b)
    state.hash.update(b)

def _write_ref(node, t):
    assert has_extrinsic(Location, node), \
            'Weak ref to unserialized: %r 0x%x' % (node, id(node))
    if isinstance(t, TVar):
        pass # Does it even make sense to check instantiations here?
    elif isinstance(t, TData):
        adt = extrinsic(FormBacking, t.data)
        assert isinstance(node, adt), "->%r is not a %s" % (node, adt)
    else:
        assert False, "%r is not a ref type" % (t,)
    loc = extrinsic(Location, node)
    a = env(Serialize).deps[loc.module]
    b = loc.index
    _write(_encode_int(a) + _encode_int(b))

def _encode_int(n):
    return unichr(n).encode('UTF-8')

def _encode_str(s):
    b = s.encode('UTF-8')
    return _encode_int(len(b)) + b

def _serialize_node(node, t):
    if isinstance(node, Structured):
        # Collect instantiations
        apps = {}
        while isinstance(t, TApply):
            target, var, arg = match(t, ("TApply(t, v, a)", tuple3))
            assert var not in apps
            apps[var] = arg
            t = target
        assert isinstance(t, TData), "%r is not a datatype" % (t,)
        adt = extrinsic(FormBacking, t.data)
        assert isinstance(node, adt), "%r is not a %s" % (node, adt)
        # Possibly write discriminator
        if len(t.data.ctors) > 1:
            ix = node._ctor_ix
            _write(_encode_int(ix))
            form = t.data.ctors[node._ctor_ix]
        else:
            form = t.data.ctors[0]
        # Dump fields
        ctor = extrinsic(FormBacking, form)
        assert isinstance(node, ctor), "%r is not a %s" % (node, ctor)
        for field in form.fields:
            sub = getattr(node, extrinsic(Name, field))
            ft = field.type
            if isinstance(ft, TWeak):
                _write_ref(sub, apps.get(ft.refType, ft.refType))
            elif isinstance(ft, TVar):
                assert ft in apps, "Can't write free %r in %r" % (ft, field)
                _serialize_node(sub, apps[ft])
            else:
                _serialize_node(sub, ft)
    elif isinstance(node, basestring):
        assert isinstance(t, TPrim) and isinstance(t.primType, PStr)
        _write(_encode_str(node))
    elif isinstance(node, int):
        assert isinstance(t, TPrim) and isinstance(t.primType, PInt)
        _write(_encode_int(node))
    elif isinstance(node, list):
        assert isinstance(t, TArray), "Unexpected array:\n%s\nfor:\n%s" % (
                node, t)
        _write(_encode_int(len(node)))
        # TODO: Check list element type for weak
        for item in node:
            _serialize_node(item, t.elemType)
    else:
        assert False, "Can't serialize %r" % (node,)

InspectState = DT('InspectState',
        ('module', '*Module'),
        ('count', int),
        ('deps', set([str])),
        )

Inspection = new_env('Inspection', InspectState)

def _inspect_node(node):
    if isinstance(node, Structured):
        assert not has_extrinsic(Location, node), "Multiply used %r" % (node,)
        state = env(Inspection)
        add_extrinsic(Location, node, Pos(state.module, state.count))
        state.count += 1
        form = node.__form__
        assert isinstance(form, Ctor)
        for field in form.fields:
            sub = getattr(node, extrinsic(Name, field))
            if isinstance(field.type, TWeak):
                # Record this ref's target digest
                if has_extrinsic(Location, sub):
                    mod = extrinsic(Location, sub).module
                    if mod != state.module:
                        state.deps.add(mod)
            else:
                _inspect_node(sub)
    elif isinstance(node, list):
        for sub in node:
            _inspect_node(sub)

def serialize(module):
    assert not has_extrinsic(ModDigest, module)
    temp = '/tmp/serialize'
    hash = sha256()
    def index():
        inspect = InspectState(module, 0, set())
        in_env(Inspection, inspect, lambda: _inspect_node(module.root))
        f = file(temp, 'wb')
        deps = []
        depmap = {module: 0}
        byDigest = [(extrinsic(ModDigest, mod), mod) for mod in inspect.deps]
        for digest, mod in sorted(byDigest):
            deps.append(digest)
            depmap[mod] = len(deps) # one-based
        state = SerialState(f, hash, 0, depmap)
        def go():
            _write(_encode_int(len(deps)))
            map_(_write, deps)
            _serialize_node(module.root, module.rootType)
        in_env(Serialize, state, go)
        f.close()
        return inspect.count
    count = index()
    hex = hash.digest().encode('hex')
    name = extrinsic(Name, module)
    add_extrinsic(ModDigest, module, hex)
    system('mv -f -- %s mods/%s' % (temp, hex))<commit_msg>Propagate TApply in native (too far though)<commit_after>from base import *
from bedrock import *
from globs import *
from hashlib import sha256
from os import system

SerialState = DT('SerialState',
        ('file', file),
        ('hash', None),
        ('index', int),
        ('deps', {Module: int}))

Serialize = new_env('Serialize', SerialState)

SAPPS = new_env('SAPPS', {TypeVar: Type})

def _write(b):
    state = env(Serialize)
    state.file.write(b)
    state.hash.update(b)

def _write_ref(node, t):
    assert has_extrinsic(Location, node), \
            'Weak ref to unserialized: %r 0x%x' % (node, id(node))
    if isinstance(t, TVar):
        pass # Does it even make sense to check instantiations here?
    elif isinstance(t, TData):
        adt = extrinsic(FormBacking, t.data)
        assert isinstance(node, adt), "->%r is not a %s" % (node, adt)
    else:
        assert False, "%r is not a ref type" % (t,)
    loc = extrinsic(Location, node)
    a = env(Serialize).deps[loc.module]
    b = loc.index
    _write(_encode_int(a) + _encode_int(b))

def _encode_int(n):
    return unichr(n).encode('UTF-8')

def _encode_str(s):
    b = s.encode('UTF-8')
    return _encode_int(len(b)) + b

def _serialize_node(node, t):
    if isinstance(node, Structured):
        # Collect instantiations
        apps = env(SAPPS).copy()
        while isinstance(t, TApply):
            target, var, arg = match(t, ("TApply(t, v, a)", tuple3))
            assert var not in apps
            apps[var] = arg
            t = target
        assert isinstance(t, TData), "%r is not a datatype" % (t,)
        adt = extrinsic(FormBacking, t.data)
        assert isinstance(node, adt), "%r is not a %s" % (node, adt)
        # Possibly write discriminator
        if len(t.data.ctors) > 1:
            ix = node._ctor_ix
            _write(_encode_int(ix))
            form = t.data.ctors[node._ctor_ix]
        else:
            form = t.data.ctors[0]
        # Dump fields
        ctor = extrinsic(FormBacking, form)
        assert isinstance(node, ctor), "%r is not a %s" % (node, ctor)
        for field in form.fields:
            sub = getattr(node, extrinsic(Name, field))
            ft = field.type
            if isinstance(ft, TVar):
                assert ft.typeVar in apps, \
                        "Can't write free %r in %r with env %r" % (
                        ft, field, apps)
                ft = apps[ft.typeVar]
                assert not isinstance(ft, TWeak), "TVar instantiated weakly?"
            if isinstance(ft, TWeak):
                _write_ref(sub, ft.refType)
            else:
                in_env(SAPPS, apps, lambda: _serialize_node(sub, ft))
    elif isinstance(node, basestring):
        assert isinstance(t, TPrim) and isinstance(t.primType, PStr)
        _write(_encode_str(node))
    elif isinstance(node, int):
        assert isinstance(t, TPrim) and isinstance(t.primType, PInt)
        _write(_encode_int(node))
    elif isinstance(node, list):
        assert isinstance(t, TArray), "Unexpected array:\n%s\nfor:\n%s" % (
                node, t)
        _write(_encode_int(len(node)))
        # TODO: Check list element type for weak
        for item in node:
            _serialize_node(item, t.elemType)
    else:
        assert False, "Can't serialize %r" % (node,)

InspectState = DT('InspectState',
        ('module', '*Module'),
        ('count', int),
        ('deps', set([str])),
        )

Inspection = new_env('Inspection', InspectState)

def _inspect_node(node):
    if isinstance(node, Structured):
        assert not has_extrinsic(Location, node), "Multiply used %r" % (node,)
        state = env(Inspection)
        add_extrinsic(Location, node, Pos(state.module, state.count))
        state.count += 1
        form = node.__form__
        assert isinstance(form, Ctor)
        for field in form.fields:
            sub = getattr(node, extrinsic(Name, field))
            if isinstance(field.type, TWeak):
                # Record this ref's target digest
                if has_extrinsic(Location, sub):
                    mod = extrinsic(Location, sub).module
                    if mod != state.module:
                        state.deps.add(mod)
            else:
                _inspect_node(sub)
    elif isinstance(node, list):
        for sub in node:
            _inspect_node(sub)

def serialize(module):
    assert not has_extrinsic(ModDigest, module)
    temp = '/tmp/serialize'
    hash = sha256()
    def index():
        inspect = InspectState(module, 0, set())
        in_env(Inspection, inspect, lambda: _inspect_node(module.root))
        f = file(temp, 'wb')
        deps = []
        depmap = {module: 0}
        byDigest = [(extrinsic(ModDigest, mod), mod) for mod in inspect.deps]
        for digest, mod in sorted(byDigest):
            deps.append(digest)
            depmap[mod] = len(deps) # one-based
        state = SerialState(f, hash, 0, depmap)
        def go():
            _write(_encode_int(len(deps)))
            map_(_write, deps)
            in_env(SAPPS, {},
                    lambda: _serialize_node(module.root, module.rootType))
        in_env(Serialize, state, go)
        f.close()
        return inspect.count
    count = index()
    hex = hash.digest().encode('hex')
    name = extrinsic(Name, module)
    add_extrinsic(ModDigest, module, hex)
    system('mv -f -- %s mods/%s' % (temp, hex))
{'fraction_non_alphanumeric': 0.12459608439460179, 'fraction_numerical': 0.003041246911233606, 'mean_word_length': 2.1734016887816647, 'pattern_counts': {'":': 0, '<': 3, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1711053', 'n_tokens_mistral': 3354, 'n_tokens_neox': 3273, 'n_words': 881}