tfhe_lbf_eval/parser.py at main · ssmiler/tfhe_lbf_eval · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
class LbfParser:
    def _strip_comments(self, line):
        # strip comments and spaces
        if (idx:= line.find("#")) >= 0:
            line = line[:idx]
        line = line.strip()
        return line

    def _get_callback(self, callback):
        def _empty_callback(*args):
            pass

        if callback is None:
            return _empty_callback
        else:
            return callback

    def __init__(self,
                 inputs_callback=None,
                 outputs_callback=None,
                 lincomb_callback=None,
                 const_callback=None,
                 bootstrap_callback=None,
                 end_callback=None):
        self.inputs_callback = self._get_callback(inputs_callback)
        self.outputs_callback = self._get_callback(outputs_callback)
        self.lincomb_callback = self._get_callback(lincomb_callback)
        self.const_callback = self._get_callback(const_callback)
        self.bootstrap_callback = self._get_callback(bootstrap_callback)
        self.end_callback = self._get_callback(end_callback)

    def parse(self, lines):
        idx = 0
        while idx < len(lines):
            line, idx = self._strip_comments(lines[idx]), idx + 1
            match line.split():
                case[".inputs", *inps]:
                    while inps[-1] == "\\":
                        line, idx = self._strip_comments(lines[idx]), idx + 1
                        inps = inps[:-1] + line.split()
                    self.inputs_callback(inps)  # TODO: support multi-line

                case[".outputs", *outs]:
                    while outs[-1] == "\\":
                        line, idx = self._strip_comments(lines[idx]), idx + 1
                        outs = outs[:-1] + line.split()
                    self.outputs_callback(outs)  # TODO: support multi-line

                case[".lincomb", *inps, out]:
                    line, idx = self._strip_comments(lines[idx]), idx + 1
                    coefs = list(map(int, line.split()))

                    const_coef = 0
                    if len(coefs) > len(inps):
                        *coefs, const_coef = coefs
                    assert(len(coefs) == len(inps)
                           ), "lincomb input and coefficient count missmatch"

                    if len(inps) > 0:
                        self.lincomb_callback(out, inps, coefs, const_coef)
                    else:
                        self.const_callback(out, const_coef)

                case[".bootstrap", inp, *outs]:
                    assert(len(outs) > 0), "bootstrap expected at least"
                    " one output"
                    tables = list()
                    for _ in range(len(outs)):
                        line, idx = self._strip_comments(lines[idx]), idx + 1
                        tables.append(list(map(int, line)))
                    self.bootstrap_callback(outs, inp, tables)

                case[".end"]:
                    self.end_callback()

                case _:
                    assert(False), f"cannot parse line no {idx}: '{line}'"


if __name__ == '__main__':
    def inputs_callback(*args):
        print(f"inputs_callback {args}")

    def outputs_callback(*args):
        print(f"outputs_callback {args}")

    def lincomb_callback(*args):
        print(f"lincomb_callback {args}")

    def const_callback(*args):
        print(f"const_callback {args}")

    def bootstrap_callback(*args):
        print(f"bootstrap_callback {args}")

    def end_callback(*args):
        print(f"end_callback {args}")

    parser = LbfParser(
        inputs_callback=inputs_callback,
        outputs_callback=outputs_callback,
        lincomb_callback=lincomb_callback,
        const_callback=const_callback,
        bootstrap_callback=bootstrap_callback,
        end_callback=end_callback)

    with open("sample.lbf", "r") as fs:
        lines = fs.readlines()
        parser.parse(lines)