recreate2.py

import sqlite3
import os
import os.path
import re
import sys
from cmd_parser import parse_line, HeapRef

def define_schema(conn):
    c = conn.cursor()

    c.execute("""
        create table Snapshot (
            id integer primary key,
            fun_call_id integer,
            start_fun_call_id integer,
            version integer,
            line_no integer,
            constraint Snapshot_fk_fun_call_id foreign key (fun_call_id)
                references FunCall(id)
        );
    """)

    c.execute("""
        create table Value (
            id integer,
            type tinyint,
            version integer,
            value,
            
            constraint Value_pk primary key (id, version)
            constraint Value_fk_type_id foreign key (type)
                references Type(id)
        );
    """)
    
    c.execute("""
        create table Member (
            container integer,
            key integer,
            value integer,
            
            constraint Member_pk primary key (container, key)
            constraint Member_fk_container foreign key (container)
                references Value(id),
            constraint Member_fk_value foreign key (value)
                references Value(id)
        );
    """)

    c.execute(""" 
        create table Type (
            id tinyint primary key,
            name
        );
    """)

    c.execute(""" 
        create table CodeFile (
            id integer primary key,
            file_path text,
            source text
        );
    """)

    c.execute("""
        create table FunCode (
            id integer primary key,
            name integer,
            code_file_id integer,
            line_no integer,
            local_varnames text,
            cell_varnames text,
            free_varnames text,

            constraint Fun_fk_code_fide_id foreign key (code_file_id)
                references CodeFile(id)
        );
    """)

    c.execute("""
        create table FunCall (
            id integer primary key,
            fun_code_id integer,
            locals integer, -- heap ID
            globals integer, -- heap ID
            cellvars text,
            freevars text,
            parent_id integer,
            
            constraint FunCall_fk_parent_id foreign key (parent_id)
                references FunCall(id)
            constraint FunCall_fk_fun_id foreign key (fun_code_id)
                references FunCode(id)
        );
    """)

    c.execute("""
        create table Error (
            id integer primary key,
            type text,
            message text,
            snapshot_id integer,

            constraint Error_fk_snapshot_id foreign key (snapshot_id)
                references Snapshot(id)
        );
    """)

    conn.commit()

class IdGenerator:
    def __init__(self):
        self.curr_id = 0
    
    def new(self):
        self.curr_id += 1
        return self.curr_id
    
    def current(self):
        return self.curr_id

class FunCode:
    def __init__(self, id, name, code_file_id,
        line_no, local_varnames, cell_varnames, free_varnames):
        self.id = id
        self.name = name
        self.code_file_id = code_file_id
        self.line_no = line_no
        self.local_varnames = local_varnames
        self.cell_varnames = cell_varnames
        self.free_varnames = free_varnames

class FunCall:
    def __init__(self, id, fun_code, 
        locals_id, globals_id, cellvars, freevars, parent):
        self.id = id
        self.fun_code = fun_code
        self.locals_id = locals_id
        self.globals_id = globals_id
        self.cellvars = cellvars
        self.freevars = freevars
        self.parent = parent

def grab_var_args(a_list):
    num_args = a_list[0]
    args = a_list[1:1 + num_args]
    rest = a_list[1 + num_args:]
    return num_args, args, rest

FROZEN_LIB_REGEX = re.compile(r"<frozen ([_a-z.A-Z0-9]+)>")
MYDIR = os.path.dirname(os.path.realpath(__file__))
def resolve_filename(filename):
    m = FROZEN_LIB_REGEX.match(filename)
    if m:
        lib = m.group(1)
        real_filename = MYDIR + "/Lib/" + "/".join(lib.split(".")) + ".py"
        return real_filename
    else:
        return filename

class HistoryRecreator:
    def __init__(self, conn, logfile):
        self.conn = conn
        self.cursor = conn.cursor()
        self.logfile = logfile
        self.file_iterator = iter(self.logfile)
        self.addr_to_id = {}
        self.type_addr_to_id = {}
        self.str_to_id = {}
        self.version = IdGenerator()
        self.value_id = IdGenerator()
        self.type_id = IdGenerator()
        self.fun_code_id = IdGenerator()
        self.code_file_id = IdGenerator()
        self.fun_call_id = IdGenerator()
        self.snapshot_id = IdGenerator()
        self.code_files = {}
        self.fun_codes = {}
        self.member_map = {} # (id, key) => value
        self.value_map = {}
        self.lists = {}
        self.stack = None
        self.curr_line_no = None
        self.buffered_lines = []
    
    def add_type(self, type_name):
        c = self.cursor
        tid = self.type_id.new()
        c.execute("insert into Type values (?, ?)", (tid, type_name))
        return tid
    
    def add_string(self, string):
        sid = self.value_id.new()
        self.insert_value(
            sid,
            self.STR_TYPE,
            self.version.current(),
            string
        )
        self.str_to_id[string] = sid
        return sid
    
    def seed_data(self):
        self.INT_TYPE = self.add_type('int')
        self.STR_TYPE = self.add_type('str')
        self.LIST_TYPE = self.add_type('list')
        self.DICT_TYPE = self.add_type('dict')
        self.MODULE_TYPE = self.add_type('module')
        self.OBJECT_TYPE = self.add_type('object')
        self.FUNCTION_TYPE = self.add_type('function')
        self.NONE_TYPE = self.add_type('none')
        self.REF_TYPE = self.add_type('<ref>')
        self.DELETED_TYPE = self.add_type('<deleted>')
        self.RET_VAL = self.add_string('<ret val>')
        self.conn.commit()
    
    def ensure_code_file_saved(self, filename):
        content = None
        if filename not in self.code_files:
            real_filename = resolve_filename(filename)
            if os.path.isfile(real_filename):
                the_file = open(real_filename, "r")
                content = the_file.read()
                the_file.close()
            code_file_id = self.code_file_id.new()
            self.cursor.execute("INSERT INTO CodeFile VALUES (?, ?, ?)", (
                code_file_id,
                filename,
                content
            ))
            self.code_files[filename] = code_file_id
    
    def rewind_basedir(self, dir):
        pass
    
    def new_code(self, heap_id, filename, name, line_no, *rest):
        num_local_vars, local_varnames, rest = grab_var_args(rest)
        num_cell_vars, cell_varnames, rest = grab_var_args(rest)
        num_free_vars, free_varnames, rest = grab_var_args(rest)
        self.ensure_code_file_saved(filename)
        fid = self.fun_code_id.new()
        code_file_id = self.code_files.get(filename)
        fun_code = FunCode(
            fid,
            name,
            code_file_id,
            line_no,
            local_varnames,
            cell_varnames,
            free_varnames
        )
        self.cursor.execute("insert into FunCode values (?, ?, ?, ?, ?, ?, ?)", (
            fun_code.id,
            fun_code.name,
            fun_code.code_file_id,
            fun_code.line_no,
            ",".join(local_varnames),
            ",".join(cell_varnames),
            ",".join(free_varnames)
        ))
        self.fun_codes[heap_id] = fun_code
    
    def visit(self, line_no):
        if line_no == -1:
            return
        self.curr_line_no = line_no
        self.save_new_snapshot()
    
    def save_new_snapshot(self):
        sid = self.snapshot_id.new()
        self.cursor.execute("insert into Snapshot values (?, ?, ?, ?, ?)", (
            sid,
            self.stack and self.stack.id,
            None,
            self.version.current(),
            self.curr_line_no
        ))
    
    def store_fast(self, idx, value):    
        self.set_local(idx, value, self.version.new())
    
    def set_local(self, idx, value, version):
        fun_code = self.stack.fun_code
        varname = fun_code.local_varnames[idx]
        if varname in self.str_to_id:
            key = self.str_to_id[varname]
        else:
            key = self.add_string(varname)
        self.set_item(self.stack.locals_id, key, value, version)
    
    def push_frame(self, code_heap_id, global_vars_id, local_vars_id, *rest):
        fun_code = self.fun_codes[code_heap_id]
        # TODO closures
        globals_id = self.get_value_id(global_vars_id)
        locals_id = self.value_id.new()
        
        cid = self.fun_call_id.new()
        parent_id = self.stack and self.stack.id or None
        
        fun_call = FunCall(
            cid,
            fun_code,
            locals_id,
            globals_id,
            None,
            None,
            self.stack
        )
        self.stack = fun_call
    
        version = self.version.new()
        self.insert_value(
            locals_id,
            self.DICT_TYPE,
            version,
            None
        )
        for i, local_var in enumerate(fun_code.local_varnames):
            value = rest[i]
            self.set_local(i, value, version)
        
        self.cursor.execute("UPDATE Snapshot set start_fun_call_id = ? where id = ?", (
            fun_call.id,
            self.snapshot_id.current()
        ))
        
        self.cursor.execute("insert into FunCall values (?, ?, ?, ?, ?, ?, ?)", (
            cid,
            fun_code.id,
            locals_id,
            globals_id,
            None, # TODO cellvars
            None, # TODO freevars
            parent_id
        ))
    
    def pop_frame(self, code_heap_id):
        self.stack = self.stack.parent
    
    def call_start(self):
        while True:
            try:
                line = next(self.file_iterator)
                command = parse_line(line)
                if command[0] == 'CALL_END':
                    # this is a call that didn't have a stack frame(FunCall) because
                    # it is a native function, we represent that by setting start_fun_call_id to 0
                    self.cursor.execute("UPDATE Snapshot set start_fun_call_id = ? where id = ?", (
                        0,
                        self.snapshot_id.current()
                    ))
                    self.save_new_snapshot()
                    break
                elif command[0] == 'PUSH_FRAME':
                    self.buffered_lines.append(line)
                    break
                else:
                    self.buffered_lines.append(line)
            except StopIteration:
                break
    
    def call_end(self):
        pass
    
    def new_string(self, addr, string):
        sid = self.get_value_id(addr)
        self.insert_value(
            sid, 
            self.STR_TYPE, 
            self.version.new(), 
            string
        )
        self.str_to_id[string] = sid
    
    def new_object(self, addr, type_name, type_addr, dict_addr=None):
        oid = self.get_value_id(addr)
        if dict_addr:
            dict_id = self.get_value_id(dict_addr)
        else:
            dict_id = None
        self.insert_value(
            oid,
            self.OBJECT_TYPE,
            self.version.new(),
            dict_id
        )
        
    def new_function(self, addr, code_heap_id):
        fun_code = self.fun_codes[code_heap_id]
        # TODO free vars
        # fun_obj = FunObject(fun_code.id)
        oid = self.get_value_id(addr)
        self.insert_value(
            oid,
            self.FUNCTION_TYPE,
            self.version.new(),
            fun_code.id
        )
    
    def new_module(self, addr):
        oid = self.get_value_id(addr)
        self.insert_value(
            oid,
            self.MODULE_TYPE,
            self.version.new(),
            None
        )
    
    def new_dict(self, addr, *entries):
        dict_id = self.get_value_id(addr)
        version = self.version.new()
        self.insert_value(
            dict_id,
            self.DICT_TYPE,
            version,
            None
        )
        for i in range(0, len(entries), 2):
            key = entries[i]
            value = entries[i + 1]
            self.set_item(dict_id, key, value, version)
            
    def dict_store_subscript(self, addr, key, value):
        dict_id = self.get_value_id(addr)
        version = self.version.new()
        self.set_item(dict_id, key, value, version)
    
    def new_list(self, addr, *items):
        list_id = self.get_value_id(addr)
        version = self.version.new()
        self.insert_value(list_id, self.LIST_TYPE, version, None)
        self.lists[list_id] = list(items)
        for i, item in enumerate(items):
            self.set_item(list_id, i, item, version, True)
    
    def list_append(self, addr, item):
        list_id = self.get_value_id(addr)
        a_list = self.lists[list_id]
        idx = len(a_list)
        self.set_item(list_id, idx, item, self.version.new())
        a_list.append(item)
    
    def list_resize_and_shift(self, addr, old_size, new_size, end_idx, range_size):
        list_id = self.get_value_id(addr)
        a_list = self.lists[list_id]
        size = len(a_list)
        version = self.version.new()
        assert (new_size - old_size) == range_size
        
        if range_size < 0:
            # shrink
            start_idx = end_idx + range_size
            num_items_to_copy = old_size - end_idx
            for i in range(num_items_to_copy):
                new_key = start_idx + i
                old_key = start_idx + i - range_size
                value = a_list[old_key]
                a_list[new_key] = value
                self.set_item(list_id, new_key, value, version)
                
            num_items_to_delete = old_size - new_size
            for i in range(num_items_to_delete):
                self.del_item(list_id, new_size + i, version)
                del a_list[new_size + i]
        else:
            # expand
            start_idx = end_idx - range_size
            num_items_to_copy = old_size - end_idx
            for i in range(num_items_to_copy):
                new_key = start_idx + i + range_size
                old_key = start_idx + i
                value = a_list[old_key]
                if len(a_list) == new_key:
                    a_list.append(value)
                else:
                    a_list[new_key] = value
                self.set_item(list_id, new_key, value, version)
    
    def list_store_index(self, addr, idx, value):
        list_id = self.get_value_id(addr)
        version = self.version.new()
        self.set_item(list_id, idx, value, version)
    
    def return_value(self, value):
        if isinstance(value, HeapRef):
            value = self.addr_to_id[value.id]
        self.set_item(
            self.stack.locals_id,
            self.RET_VAL,
            value,
            self.version.new()
        )
        self.save_new_snapshot()
    
    def set_item(self, dict_id, key, value, version):
        if isinstance(key, HeapRef):
            key = self.addr_to_id[key.id]
        if isinstance(value, HeapRef):
            value = self.addr_to_id[value.id]
            ref = self.get_ref_id(dict_id, key, version)
            self.insert_value(
                ref,
                self.REF_TYPE,
                version,
                value
            )
        else:
            ref = self.get_ref_id(dict_id, key, version)
            type_id = self.get_type_by_value(value)
            self.insert_value(
                ref,
                type_id,
                version,
                value
            )
    
    def del_item(self, container_id, key, version):
        if isinstance(key, HeapRef):
            key = self.addr_to_id[key.id]
        ref = self.get_ref_id(container_id, key, version)
        self.insert_value(
            ref,
            self.DELETED_TYPE,
            version,
            None
        )
    
    def get_type_by_value(self, value):
        if value is None:
            return self.NONE_TYPE
        if isinstance(value, int):
            return self.INT_TYPE
        else:
            raise Exception("Unsupported type: %r" % type(value))
        
    def get_ref_id(self, dict_id, key, version):
        if (dict_id, key) in self.member_map:
            return self.member_map[(dict_id, key)]
        else:
            ref = self.value_id.new()
            self.cursor.execute("insert into Member values (?, ?, ?)", (
                dict_id,
                key,
                ref
            ))
            self.member_map[(dict_id, key)] = ref
            return ref

    def insert_value(self, oid, type_id, version, value):
        self.cursor.execute("insert into Value values (?, ?, ?, ?)", (
            oid,
            type_id,
            version,
            value
        ))
        
    def get_type_id(self, type_name, type_addr):
        if type_addr in self.type_addr_to_id:
            type_id = self.type_addr_to_id[type_addr]
        else:
            type_id = self.add_type(type_name)
            self.type_addr_to_id[type_addr] = type_id
        return type_id
        
    def get_value_id(self, addr):
        if addr in self.addr_to_id:
            return self.addr_to_id[addr]
        else:
            id = self.value_id.new()
            self.addr_to_id[addr] = id
            return id
    
    def process_command(self, command):
        meth_name = command[0].lower()
        args = command[1]
        if hasattr(self, meth_name):
            getattr(self, meth_name)(*args)
        else:
            print("Warning no matching method for %s on line %d" % (meth_name, log_line_no))
    
    def recreate(self):
        self.seed_data()
        log_line_no = 0
        try:
            while True:
                if len(self.buffered_lines) > 0:
                    line = self.buffered_lines.pop(0)
                else:
                    try:
                        line = next(self.file_iterator)
                    except StopIteration:
                        break
                log_line_no += 1
                command = parse_line(line)
                self.process_command(command)
                if log_line_no % 500 == 0:
                    self.conn.commit()
        except StopIteration:
            
            pass
        except Exception as e:
            print("ERROR on line %d: %s" % (log_line_no, line))
            raise e
        self.conn.commit()
        
def main():
    if len(sys.argv) < 2:
        print("Please provide a .rewind file.")
        return
    filename = sys.argv[1]

    if not filename.endswith(".rewind"):
        print("Please provide a .rewind file.")
        return
    
    sqlite_filename = re.sub("\.rewind$", ".test.sqlite", filename)
    if os.path.isfile(sqlite_filename):
        os.remove(sqlite_filename)
    conn = sqlite3.connect(sqlite_filename)
    define_schema(conn)
    logfile = open(filename, "r")
    recreator = HistoryRecreator(conn, logfile)
    recreator.recreate()

if __name__ == "__main__":
    main()