1import claripy
2import nose
3
4def test_simple_merging():
5    yield raw_simple_merging, claripy.Solver
6    yield raw_simple_merging, claripy.SolverHybrid
7    yield raw_simple_merging, claripy.SolverComposite
8
9def raw_simple_merging(solver_type):
10    s1 = solver_type()
11    s2 = solver_type()
12    w = claripy.BVS("w", 8)
13    x = claripy.BVS("x", 8)
14    y = claripy.BVS("y", 8)
15    z = claripy.BVS("z", 8)
16    m = claripy.BVS("m", 8)
17
18    s1.add([x == 1, y == 10])
19    s2.add([x == 2, z == 20, w == 5])
20    _, sm = s1.merge([s2], [ m == 0, m == 1 ])
21
22    nose.tools.assert_equal(s1.eval(x, 1), (1,))
23    nose.tools.assert_equal(s2.eval(x, 1), (2,))
24
25    sm1 = sm.branch()
26    sm1.add(x == 1)
27    nose.tools.assert_equal(sm1.eval(x, 1), (1,))
28    nose.tools.assert_equal(sm1.eval(y, 1), (10,))
29    nose.tools.assert_equal(sm1.eval(z, 1), (0,))
30    nose.tools.assert_equal(sm1.eval(w, 1), (0,))
31
32    sm2 = sm.branch()
33    sm2.add(x == 2)
34    nose.tools.assert_equal(sm2.eval(x, 1), (2,))
35    nose.tools.assert_equal(sm2.eval(y, 1), (0,))
36    nose.tools.assert_equal(sm2.eval(z, 1), (20,))
37    nose.tools.assert_equal(sm2.eval(w, 1), (5,))
38
39    sm1 = sm.branch()
40    sm1.add(m == 0)
41    nose.tools.assert_equal(sm1.eval(x, 1), (1,))
42    nose.tools.assert_equal(sm1.eval(y, 1), (10,))
43    nose.tools.assert_equal(sm1.eval(z, 1), (0,))
44    nose.tools.assert_equal(sm1.eval(w, 1), (0,))
45
46    sm2 = sm.branch()
47    sm2.add(m == 1)
48    nose.tools.assert_equal(sm2.eval(x, 1), (2,))
49    nose.tools.assert_equal(sm2.eval(y, 1), (0,))
50    nose.tools.assert_equal(sm2.eval(z, 1), (20,))
51    nose.tools.assert_equal(sm2.eval(w, 1), (5,))
52
53    m2 = claripy.BVS("m2", 32)
54    _, smm = sm1.merge([sm2], [ m2 == 0, m2 == 1 ])
55
56    smm_1 = smm.branch()
57    smm_1.add(x == 1)
58    nose.tools.assert_equal(smm_1.eval(x, 1), (1,))
59    nose.tools.assert_equal(smm_1.eval(y, 1), (10,))
60    nose.tools.assert_equal(smm_1.eval(z, 1), (0,))
61    nose.tools.assert_equal(smm_1.eval(w, 1), (0,))
62
63    smm_2 = smm.branch()
64    smm_2.add(m == 1)
65    nose.tools.assert_equal(smm_2.eval(x, 1), (2,))
66    nose.tools.assert_equal(smm_2.eval(y, 1), (0,))
67    nose.tools.assert_equal(smm_2.eval(z, 1), (20,))
68    nose.tools.assert_equal(smm_2.eval(w, 1), (5,))
69
70    so = solver_type()
71    so.add(w == 0)
72
73    sa = so.branch()
74    sb = so.branch()
75    sa.add(x == 1)
76    sb.add(x == 2)
77    _, sm = sa.merge([sb], [ m == 0, m == 1 ])
78
79    smc = sm.branch()
80    smd = sm.branch()
81    smc.add(y == 3)
82    smd.add(y == 4)
83
84    _, smm = smc.merge([smd], [ m2 == 0, m2 == 1 ])
85    wxy = claripy.Concat(w, x, y)
86
87    smm_1 = smm.branch()
88    smm_1.add(wxy == 0x000103)
89    nose.tools.assert_true(smm_1.satisfiable())
90
91    smm_1 = smm.branch()
92    smm_1.add(wxy == 0x000104)
93    nose.tools.assert_true(smm_1.satisfiable())
94
95    smm_1 = smm.branch()
96    smm_1.add(wxy == 0x000203)
97    nose.tools.assert_true(smm_1.satisfiable())
98
99    smm_1 = smm.branch()
100    smm_1.add(wxy == 0x000204)
101    nose.tools.assert_true(smm_1.satisfiable())
102
103    smm_1 = smm.branch()
104    smm_1.add(wxy != 0x000103)
105    smm_1.add(wxy != 0x000104)
106    smm_1.add(wxy != 0x000203)
107    smm_1.add(wxy != 0x000204)
108    nose.tools.assert_false(smm_1.satisfiable())
109
110if __name__ == '__main__':
111    for func, param in test_simple_merging():
112        func(param)
113