from z3 import *

#------------------------------------------------------
# A first example.  Reasoning with Z3
#------------------------------------------------------

# first let's make a sort.  we can think of this as introducing a class
Entity = DeclareSort('Entity')


# human-func is a function from (Entity) to Bool
human_func = Function('human_func', Entity, BoolSort())
mortal_func = Function('mortal_func', Entity, BoolSort())

# let's introduce someone interesting!
socrates = Const('socrates', Entity)

# ok, now let's make a solver, the thing that'll do our work for us!
s = Solver()

# we're going to want to use a variable in a forall
# free variables used in forall have to be declared Const
x = Const('x', Entity)

# all humans are mortal
axiom1 = ForAll([x], Implies(human_func(x), mortal_func(x)))

# Socrates is a human
axiom2 = human_func(socrates)

# let's tell our solver we know these things
s.add([axiom1, axiom2])

# and let's ask the solver if these things are coherent
print("coherent so far?", s.check())

# we're going for: all humans are mortal; Socrates is a human; therefore Socrates is mortal
# but we don't just want to see if it's *possible* to give Socrates a type that makes him
# mortal.  we want to make sure it's impossible for him to be immortal
# (remember the difference between satisfiability and validity we talked about last class;
# we're looking for validity)
# so...

s.add(Not(mortal_func(socrates)))
print("possible that Socrates is not mortal?", s.check())


#------------------------------------------------------
# Example 2.  Filling a hole
#------------------------------------------------------

x = BitVec('x', 8)

# x * 2 is slow, but maybe you remember there's a cute bit vector trick
# that does the same thing but faster, with the form x << ??
slow_expr = x * 2
# remember ?? from Rosette and our discussions of Sketch?
# here we'll introduce a variable h that represents our hole
h = BitVec('h', 8)

fast_expr = x << h

goal = ForAll([x], slow_expr == fast_expr)

s = Solver()
s.add(goal)
s.check()
print("how should we fill this hole?", s.model()) # we've made a subexpression!

#------------------------------------------------------
# Example 3.  Just a little intro to Z3 Ifs
#------------------------------------------------------

x = Int('x')
s = Solver()
# ifs are going to be a little different for constraints...
s.add(If(True, x==3, x==10)) 
s.check()
print("what should x be, based on our If?", s.model())