dawg.py

# Original Algorithm:
# By Steve Hanov, 2011. Released to the public domain.
# Please see http://stevehanov.ca/blog/index.php?id=115 for the accompanying article.
#
# Adapted for PyPy/CPython by Carl Friedrich Bolz-Tereick
#
# Based on Daciuk, Jan, et al. "Incremental construction of minimal acyclic finite-state automata."
# Computational linguistics 26.1 (2000): 3-16.
#
# Updated 2014 to use DAWG as a mapping; see
# Kowaltowski, T.; CL. Lucchesi (1993), "Applications of finite automata representing large vocabularies",
# Software-Practice and Experience 1993

fromcollectionsimportdefaultdict
fromfunctoolsimportcached_property


# This class represents a node in the directed acyclic word graph (DAWG). It
# has a list of edges to other nodes. It has functions for testing whether it
# is equivalent to another node. Nodes are equivalent if they have identical
# edges, and each identical edge leads to identical states. The __hash__ and
# __eq__ functions allow it to be used as a key in a python dictionary.


classDawgNode:

def__init__(self, dawg):
self.id=dawg.next_id
dawg.next_id+=1
self.final=False
self.edges= {}

self.linear_edges=None# later: list of (string, next_state)

def__str__(self):
ifself.final:
arr= ["1"]
else:
arr= ["0"]

for (label, node) insorted(self.edges.items()):
arr.append(label)
arr.append(str(node.id))

return"_".join(arr)
__repr__=__str__

def_as_tuple(self):
edges=sorted(self.edges.items())
edge_tuple=tuple((label, node.id) forlabel, nodeinedges)
return (self.final, edge_tuple)

def__hash__(self):
returnhash(self._as_tuple())

def__eq__(self, other):
returnself._as_tuple() ==other._as_tuple()

@cached_property
defnum_reachable_linear(self):
# returns the number of different paths to final nodes reachable from
# this one

count=0
# staying at self counts as a path if self is final
ifself.final:
count+=1
forlabel, nodeinself.linear_edges:
count+=node.num_reachable_linear

returncount


classDawg:
def__init__(self):
self.previous_word=""
self.next_id=0
self.root=DawgNode(self)

# Here is a list of nodes that have not been checked for duplication.
self.unchecked_nodes= []

# To deduplicate, maintain a dictionary with
# minimized_nodes[canonical_node] is canonical_node.
# Based on __hash__ and __eq__, minimized_nodes[n] is the
# canonical node equal to n.
# In other words, self.minimized_nodes[x] == x for all nodes found in
# the dict.
self.minimized_nodes= {}

# word: value mapping
self.data= {}
# value: word mapping
self.inverse= {}

definsert(self, word, value):
ifnotall(0<=ord(c) <128forcinword):
raiseValueError("Use 7-bit ASCII characters only")
ifword<=self.previous_word:
raiseValueError("Error: Words must be inserted in alphabetical order.")
ifvalueinself.inverse:
raiseValueError(f"value {value} is duplicate, got it for word {self.inverse[value]} and now {word}")

# find common prefix between word and previous word
common_prefix=0
foriinrange(min(len(word), len(self.previous_word))):
ifword[i] !=self.previous_word[i]:
break
common_prefix+=1

# Check the unchecked_nodes for redundant nodes, proceeding from last
# one down to the common prefix size. Then truncate the list at that
# point.
self._minimize(common_prefix)

self.data[word] =value
self.inverse[value] =word

# add the suffix, starting from the correct node mid-way through the
# graph
iflen(self.unchecked_nodes) ==0:
node=self.root
else:
node=self.unchecked_nodes[-1][2]

forletterinword[common_prefix:]:
next_node=DawgNode(self)
node.edges[letter] =next_node
self.unchecked_nodes.append((node, letter, next_node))
node=next_node

node.final=True
self.previous_word=word

deffinish(self):
ifnotself.data:
raiseValueError("need at least one word in the dawg")
# minimize all unchecked_nodes
self._minimize(0)

self._linearize_edges()

topoorder, linear_data, inverse=self._topological_order()
returnself.compute_packed(topoorder), linear_data, inverse

def_minimize(self, down_to):
# proceed from the leaf up to a certain point
foriinrange(len(self.unchecked_nodes) -1, down_to-1, -1):
 (parent, letter, child) =self.unchecked_nodes[i]
ifchildinself.minimized_nodes:
# replace the child with the previously encountered one
parent.edges[letter] =self.minimized_nodes[child]
else:
# add the state to the minimized nodes.
self.minimized_nodes[child] =child
self.unchecked_nodes.pop()

def_lookup(self, word):
""" Return an integer 0 <= k < number of strings in dawg
 where word is the kth successful traversal of the dawg. """
node=self.root
skipped=0# keep track of number of final nodes that we skipped
index=0
whileindex<len(word):
forlabel, childinnode.linear_edges:
ifword[index] ==label[0]:
ifword[index:index+len(label)] ==label:
ifnode.final:
skipped+=1
index+=len(label)
node=child
break
else:
returnNone
skipped+=child.num_reachable_linear
else:
returnNone
returnskipped

defenum_all_nodes(self):
stack= [self.root]
done=set()
whilestack:
node=stack.pop()
ifnode.idindone:
continue
yieldnode
done.add(node.id)
forlabel, childinsorted(node.edges.items()):
stack.append(child)

defprettyprint(self):
fornodeinsorted(self.enum_all_nodes(), key=lambdae: e.id):
s_final=" final"ifnode.finalelse""
print(f"{node.id}: ({node}) {s_final}")
forlabel, childinsorted(node.edges.items()):
print(f" {label} goto {child.id}")

def_inverse_lookup(self, number):
assert0, "not working in the current form, but keep it as the pure python version of compact lookup"
result= []
node=self.root
while1:
ifnode.final:
ifpos==0:
return"".join(result)
pos-=1
forlabel, childinsorted(node.edges.items()):
nextpos=pos-child.num_reachable_linear
ifnextpos<0:
result.append(label)
node=child
break
else:
pos=nextpos
else:
assert0

def_linearize_edges(self):
# compute "linear" edges. the idea is that long chains of edges without
# any of the intermediate states being final or any extra incoming or
# outgoing edges can be represented by having removing them, and
# instead using longer strings as edge labels (instead of single
# characters)
incoming=defaultdict(list)
nodes=sorted(self.enum_all_nodes(), key=lambdae: e.id)
fornodeinnodes:
forlabel, childinsorted(node.edges.items()):
incoming[child].append(node)
fornodeinnodes:
node.linear_edges= []
forlabel, childinsorted(node.edges.items()):
s= [label]
whilelen(child.edges) ==1andlen(incoming[child]) ==1andnotchild.final:
 (c, child), =child.edges.items()
s.append(c)
node.linear_edges.append((''.join(s), child))

def_topological_order(self):
# compute reachable linear nodes, and the set of incoming edges for each node
order= []
stack= [self.root]
seen=set()
whilestack:
# depth first traversal
node=stack.pop()
ifnode.idinseen:
continue
seen.add(node.id)
order.append(node)
forlabel, childinnode.linear_edges:
stack.append(child)

# do a (slightly bad) topological sort
incoming=defaultdict(set)
fornodeinorder:
forlabel, childinnode.linear_edges:
incoming[child].add((label, node))
no_incoming= [order[0]]
topoorder= []
positions= {}
whileno_incoming:
node=no_incoming.pop()
topoorder.append(node)
positions[node] =len(topoorder)
# use "reversed" to make sure that the linear_edges get reorderd
# from their alphabetical order as little as necessary (no_incoming
# is LIFO)
forlabel, childinreversed(node.linear_edges):
incoming[child].discard((label, node))
ifnotincoming[child]:
no_incoming.append(child)
delincoming[child]
# check result
assertset(topoorder) ==set(order)
assertlen(set(topoorder)) ==len(topoorder)

fornodeinorder:
node.linear_edges.sort(key=lambdaelement: positions[element[1]])

fornodeinorder:
forlabel, childinnode.linear_edges:
assertpositions[child] >positions[node]
# number the nodes. afterwards every input string in the set has a
# unique number in the 0 <= number < len(data). We then put the data in
# self.data into a linear list using these numbers as indexes.
topoorder[0].num_reachable_linear
linear_data= [None] *len(self.data)
inverse= {} # maps value back to index
forword, valueinself.data.items():
index=self._lookup(word)
linear_data[index] =value
inverse[value] =index

returntopoorder, linear_data, inverse

defcompute_packed(self, order):
defcompute_chunk(node, offsets):
""" compute the packed node/edge data for a node. result is a
 list of bytes as long as order. the jump distance calculations use
 the offsets dictionary to know where in the final big output
 bytestring the individual nodes will end up. """
result=bytearray()
offset=offsets[node]
encode_varint_unsigned(number_add_bits(node.num_reachable_linear, node.final), result)
iflen(node.linear_edges) ==0:
assertnode.final
encode_varint_unsigned(0, result) # add a 0 saying "done"
prev_child_offset=offset+len(result)
foredgeindex, (label, targetnode) inenumerate(node.linear_edges):
label=label.encode('ascii')
child_offset=offsets[targetnode]
child_offset_difference=child_offset-prev_child_offset

info=number_add_bits(child_offset_difference, len(label) ==1, edgeindex==len(node.linear_edges) -1)
ifedgeindex==0:
assertinfo!=0
encode_varint_unsigned(info, result)
prev_child_offset=child_offset
iflen(label) >1:
encode_varint_unsigned(len(label), result)
result.extend(label)
returnresult

defcompute_new_offsets(chunks, offsets):
""" Given a list of chunks, compute the new offsets (by adding the
 chunk lengths together). Also check if we cannot shrink the output
 further because none of the node offsets are smaller now. if that's
 the case return None. """
new_offsets= {}
curr_offset=0
should_continue=False
fornode, resultinzip(order, chunks):
ifcurr_offset<offsets[node]:
# the new offset is below the current assumption, this
# means we can shrink the output more
should_continue=True
new_offsets[node] =curr_offset
curr_offset+=len(result)
ifnotshould_continue:
returnNone
returnnew_offsets

# assign initial offsets to every node
offsets= {}
fori, nodeinenumerate(order):
# we don't know position of the edge yet, just use something big as
# the starting position. we'll have to do further iterations anyway,
# but the size is at least a lower limit then
offsets[node] =i*2**30


# due to the variable integer width encoding of edge targets we need to
# run this to fixpoint. in the process we shrink the output more and
# more until we can't any more. at any point we can stop and use the
# output, but we might need padding zero bytes when joining the chunks
# to have the correct jump distances
last_offsets=None
while1:
chunks= [compute_chunk(node, offsets) fornodeinorder]
last_offsets=offsets
offsets=compute_new_offsets(chunks, offsets)
ifoffsetsisNone: # couldn't shrink
break

# build the final packed string
total_result=bytearray()
fornode, resultinzip(order, chunks):
node_offset=last_offsets[node]
ifnode_offset>len(total_result):
# need to pad to get the offsets correct
padding=b"\x00"* (node_offset-len(total_result))
total_result.extend(padding)
assertnode_offset==len(total_result)
total_result.extend(result)
returnbytes(total_result)


# ______________________________________________________________________
# the following functions operate on the packed representation

defnumber_add_bits(x, *bits):
forbitinbits:
assertbit==0orbit==1
x= (x<<1) |bit
returnx

defencode_varint_unsigned(i, res):
# https://en.wikipedia.org/wiki/LEB128 unsigned variant
more=True
startlen=len(res)
ifi<0:
raiseValueError("only positive numbers supported", i)
whilemore:
lowest7bits=i&0b1111111
i>>=7
ifi==0:
more=False
else:
lowest7bits|=0b10000000
res.append(lowest7bits)
returnlen(res) -startlen

defnumber_split_bits(x, n, acc=()):
ifn==1:
returnx>>1, x&1
ifn==2:
returnx>>2, (x>>1) &1, x&1
assert0, "implement me!"

defdecode_varint_unsigned(b, index=0):
res=0
shift=0
whileTrue:
byte=b[index]
res=res| ((byte&0b1111111) <<shift)
index+=1
shift+=7
ifnot (byte&0b10000000):
returnres, index

defdecode_node(packed, node):
x, node=decode_varint_unsigned(packed, node)
node_count, final=number_split_bits(x, 1)
returnnode_count, final, node

defdecode_edge(packed, edgeindex, prev_child_offset, offset):
x, offset=decode_varint_unsigned(packed, offset)
ifx==0andedgeindex==0:
raiseKeyError# trying to decode past a final node
child_offset_difference, len1, last_edge=number_split_bits(x, 2)
child_offset=prev_child_offset+child_offset_difference
iflen1:
size=1
else:
size, offset=decode_varint_unsigned(packed, offset)
returnchild_offset, last_edge, size, offset

def_match_edge(packed, s, size, node_offset, stringpos):
ifsize>1andstringpos+size>len(s):
# past the end of the string, can't match
returnFalse
foriinrange(size):
ifpacked[node_offset+i] !=s[stringpos+i]:
# if a subsequent char of an edge doesn't match, the word isn't in
# the dawg
ifi>0:
raiseKeyError
returnFalse
returnTrue

deflookup(packed, data, s):
returndata[_lookup(packed, s)]

def_lookup(packed, s):
stringpos=0
node_offset=0
skipped=0# keep track of number of final nodes that we skipped
false=False
whilestringpos<len(s):
#print(f"{node_offset=} {stringpos=}")
_, final, edge_offset=decode_node(packed, node_offset)
prev_child_offset=edge_offset
edgeindex=0
while1:
child_offset, last_edge, size, edgelabel_chars_offset=decode_edge(packed, edgeindex, prev_child_offset, edge_offset)
#print(f" {edge_offset=} {child_offset=} {last_edge=} {size=} {edgelabel_chars_offset=}")
edgeindex+=1
prev_child_offset=child_offset
if_match_edge(packed, s, size, edgelabel_chars_offset, stringpos):
# match
iffinal:
skipped+=1
stringpos+=size
node_offset=child_offset
break
iflast_edge:
raiseKeyError
descendant_count, _, _=decode_node(packed, child_offset)
skipped+=descendant_count
edge_offset=edgelabel_chars_offset+size
_, final, _=decode_node(packed, node_offset)
iffinal:
returnskipped
raiseKeyError

definverse_lookup(packed, inverse, x):
pos=inverse[x]
return_inverse_lookup(packed, pos)

def_inverse_lookup(packed, pos):
result=bytearray()
node_offset=0
while1:
node_count, final, edge_offset=decode_node(packed, node_offset)
iffinal:
ifpos==0:
returnbytes(result)
pos-=1
prev_child_offset=edge_offset
edgeindex=0
while1:
child_offset, last_edge, size, edgelabel_chars_offset=decode_edge(packed, edgeindex, prev_child_offset, edge_offset)
edgeindex+=1
prev_child_offset=child_offset
descendant_count, _, _=decode_node(packed, child_offset)
nextpos=pos-descendant_count
ifnextpos<0:
assertedgelabel_chars_offset>=0
result.extend(packed[edgelabel_chars_offset: edgelabel_chars_offset+size])
node_offset=child_offset
break
elifnotlast_edge:
pos=nextpos
edge_offset=edgelabel_chars_offset+size
else:
raiseKeyError
else:
raiseKeyError


defbuild_compression_dawg(ucdata):
d=Dawg()
ucdata.sort()
forname, valueinucdata:
d.insert(name, value)
packed, pos_to_code, reversedict=d.finish()
print("size of dawg [KiB]", round(len(packed) /1024, 2))
# check that lookup and inverse_lookup work correctly on the input data
forname, valueinucdata:
assertlookup(packed, pos_to_code, name.encode('ascii')) ==value
assertinverse_lookup(packed, reversedict, value) ==name.encode('ascii')
returnpacked, pos_to_code