mkmarkov

#!/usr/bin/python
from sys import stdin
try:
	from pydot import Dot, Node, Edge, Graph
	write_dot = True
except ImportError:
	write_dot = False

def do_node(dot, visitmap, node):
	if visitmap.has_key(node):
		return visitmap[node]

	if node.uri == None:
		n = Node("%r"%node)
		n.set_label("ENTRY")
		n.set_color("#b0ffb0")
	else:
		n = Node("%s"%node)
		n.set_label("%s"%node)
	n.set_style("filled")
	dot.add_node(n)
	visitmap[node] = n

	for (u, w) in node.links:
		c = do_node(dot, visitmap, u)
		e = Edge(n, c)
		e.set_label("Weight = %u"%w)
		dot.add_edge(e)

	for sub in node.sub_uris:
		sn = Node("t_" + sub)
		sn.set_label("ANCILLARY\\n%s"%sub)
		sn.set_color("#ffb0b0")
		sn.set_style("filled")
		sn.set_shape("rectangle")
		dot.add_node(sn)
		e = Edge(n, sn)
		dot.add_edge(e)

	return n

def write_graphviz_model(fn, initial):
	dot = Dot(graph_name="\"HTTP client markov model\"")
	dot.set_label("HTTP client markov model")
	dot.set_rankdir("LR")
	visited = {}
	do_node(dot, visited, initial)
	dot.write(fn)
	print "written: %s"%fn

def do_get_nodelist(visitmap, list, node):
	if visitmap.has_key(node):
		return
	visitmap[node] = True
	list.append(node)
	for (u, w) in node.links:
		do_get_nodelist(visitmap, list, u)

def get_nodelist(initial):
	visitmap = {}
	list = []
	do_get_nodelist(visitmap, list, initial)
	return list

class HTTPObject:
	def __init__(self, uri):
		if uri != None:
			assert(uri[0] == '/')
		self.uri = uri
		self.sub_uris = []
		self.__links = []
		self.links = []
	def add_sub(self, uri):
		self.sub_uris.append(uri)
	def add_link(self, uri, weight):
		self.__links.append((uri, weight))
	def set_name(self, name):
		self.name = name
	def resolve_links(self, objmap):
		self.links = []
		for (u, w) in self.__links:
			if objmap.has_key(u):
				self.links.append((objmap[u], w))
			else:
				print "Non-specified terminal: %s"%u
				self.links.append((HTTPObject(u), w))
		del self.__links
	def __str__(self):
		return self.uri

def declare_node(f, n):
	f.write("static const struct mnode %s;\n"%n.name)

def define_edges(f, n):
	f.write("static const struct medge %s_edges[] = {\n"%n.name)
	prob_max = 0
	for (obj, weight)in n.links:
		prob_max += weight
		f.write("\t{ .e_prob_max = %u, .e_node = &%s},\n"%(\
			prob_max, obj.name))
	f.write("};\n\n")

def num_bits(num):
	i = 0
	while True:
		if (1 << i) >= num:
			return i
		i += 1

def define_node(f, n):
	f.write("static const struct mnode %s = {\n"%n.name)
	if n.uri:
		f.write("\t.n_uri = { .v_ptr = (uint8_t *)\"%s\", "%n.uri +
			".v_len = %u },\n"%len(n.uri))
	if len(n.sub_uris):
		f.write("\t.n_num_sub = %u,\n"%len(n.sub_uris))
		f.write("\t.n_ancillary = (struct ro_vec []){\n")
		for sub in n.sub_uris:
			f.write("\t\t{ .v_ptr = (uint8_t *)\"%s\", "%sub)
			f.write(".v_len = %u },\n"%len(sub))
		f.write("\t},\n")
	if len(n.links):
		f.write("\t.n_num_edges = %u,\n"%len(n.links))
		f.write("\t.n_edge_prob_bits = %u,\n"%\
			num_bits(sum(map(lambda x:x[1], n.links))))
		f.write("\t.n_edge_prob_max = %u,\n"%\
			sum(map(lambda x:x[1], n.links)))
		f.write("\t.n_edges = %s_edges,\n"%n.name)
	f.write("};\n\n")

if __name__ == '__main__':
	cur_obj = None
	objmap = {}
	initial = HTTPObject(None)
	for line in stdin:
		line = line.rstrip("\r\n")
		if len(line) == 0 or line[0] == '#':
			continue
		if line[0] == '\t':
			line = line[1:]
			if len(line) > 2 and line[:2] == '->':
				toks = line.split(None, 1)
				uri = toks[0][2:]
				if len(toks) > 1:
					weight = int(toks[1])
				else:
					weight = 1
				cur_obj.add_link(uri, weight)
			else:
				cur_obj.add_sub(line)
		else:
			toks = line.split(None, 1)
			uri = toks[0]
			if len(toks) > 1:
				weight = int(toks[1])
			else:
				weight = 1
			initial.add_link(uri, weight)
			cur_obj = HTTPObject(uri)
			objmap[cur_obj.uri] = cur_obj

	for x in objmap.values():
		x.resolve_links(objmap)
	initial.resolve_links(objmap)
	del objmap

	if write_dot:
		write_graphviz_model("markov.dot", initial)
	
	f = open("markov.c", 'w')
	f.write("/* automatically generated by mkmarkov */\n")
	f.write("#include <httprape.h>\n\n")
	nodelist = get_nodelist(initial)
	print "In-order traversal of graph:"
	i = 0
	for x in nodelist:
		if x.uri:
			x.set_name("node%u"%i)
			i += 1
		else:
			x.set_name("_markov_root")
		declare_node(f, x)
	f.write("\n")

	for x in nodelist:
		if len(x.links):
			define_edges(f, x)
	for x in nodelist:
		define_node(f, x)
	f.write("const struct mnode *markov_root = &_markov_root;\n")