import py

import sys

import math

from rpython.rlib.rarithmetic import intmask, ovfcheck

from rpython.rlib.rarithmetic import r_uint, LONG_BIT

LOG2 = math.log(2)

NBITS = int(math.log(sys.maxint) / LOG2) + 2

# XXX should optimize the numbers

NEW_NODE_WHEN_LENGTH = 32

CONVERT_WHEN_SMALLER = 8

MAX_DEPTH = 32 # maybe should be smaller

CONCATENATE_WHEN_MULTIPLYING = 128

HIGHEST_BIT_SET = intmask(1L << (NBITS - 1))

def find_fib_index(l):

if l == 0:

return -1

a, b = 1, 2

i = 0

while 1:

if a <= l < b:

return i

a, b = b, a + b

i += 1

mul_by_1000003_pow_table = [intmask(pow(1000003, 2**_i, 2**LONG_BIT))

for _i in range(LONG_BIT)]

def masked_mul_by_1000003_pow(a, b):

"""Computes intmask(a * (1000003**b))."""

index = 0

b = r_uint(b)

while b:

if b & 1:

a = intmask(a * mul_by_1000003_pow_table[index])

b >>= 1

index += 1

return a

class GlobalRopeInfo(object):

"""contains data that is "global" to the whole string, e.g. requires

an iteration over the whole string"""

def __init__(self):

self.charbitmask = 0

self.hash = 0

self.is_bytestring = False

self.is_ascii = False

def combine(self, other, rightlength):

result = GlobalRopeInfo()

result.charbitmask = self.charbitmask | other.charbitmask

h1 = self.hash

h2 = other.hash

x = intmask(h2 + masked_mul_by_1000003_pow(h1, rightlength))

x |= HIGHEST_BIT_SET

result.hash = x

result.is_ascii = self.is_ascii and other.is_ascii

result.is_bytestring = self.is_bytestring and other.is_bytestring

return result

class StringNode(object):

_additional_info = None

def additional_info(self):

addinfo = self._additional_info

if addinfo is None:

return self.compute_additional_info()

return addinfo

def compute_additional_info(self):

raise NotImplementedError("base class")

def length(self):

raise NotImplementedError("base class")

def is_ascii(self):

raise NotImplementedError("base class")

def is_bytestring(self):

raise NotImplementedError("base class")

def depth(self):

return 0

def hash_part(self):

raise NotImplementedError("base class")

def charbitmask(self):

raise NotImplementedError("base class")

def check_balanced(self):

return True

def getchar(self, index):

raise NotImplementedError("abstract base class")

def getunichar(self, index):

raise NotImplementedError("abstract base class")

def getint(self, index):

raise NotImplementedError("abstract base class")

def getrope(self, index):

raise NotImplementedError("abstract base class")

def getslice(self, start, stop):

raise NotImplementedError("abstract base class")

def can_contain_int(self, value):

return True #conservative default

def view(self):

view([self])

def rebalance(self):

return self

def flatten_string(self):

raise NotImplementedError("abstract base class")

def flatten_unicode(self):

raise NotImplementedError("abstract base class")

def _concat(self, other):

raise NotImplementedError("abstract base class")

def __add__(self, other):

return concatenate(self, other)

def _cleanup_(self):

self.additional_info()

class LiteralNode(StringNode):

def find_int(self, what, start, stop):

raise NotImplementedError("abstract base class")

class LiteralStringNode(LiteralNode):

def __init__(self, s):

assert isinstance(s, str)

self.s = s

def length(self):

return len(self.s)

def is_ascii(self):

return self.additional_info().is_ascii

def is_bytestring(self):

return True

def flatten_string(self):

return self.s

def flatten_unicode(self):

return self.s.decode('latin-1')

def hash_part(self):

return self.additional_info().hash

def compute_additional_info(self):

additional_info = GlobalRopeInfo()

is_ascii = True

charbitmask = 0

partial_hash = 0

for c in self.s:

ordc = ord(c)

partial_hash = (1000003*partial_hash) + ord(c)

if ordc >= 128:

is_ascii = False

charbitmask |= intmask(1 << (ordc & 0x1F))

partial_hash = intmask(partial_hash)

partial_hash |= HIGHEST_BIT_SET

additional_info.hash = partial_hash

additional_info.is_ascii = is_ascii

additional_info.charbitmask = charbitmask

additional_info.is_bytestring = True

self._additional_info = additional_info

return additional_info

def charbitmask(self):

return self.additional_info().charbitmask

def getchar(self, index):

return self.s[index]

def getunichar(self, index):

return unichr(ord(self.s[index]))

def getint(self, index):

return ord(self.s[index])

def getrope(self, index):

return LiteralStringNode.PREBUILT[ord(self.s[index])]

def can_contain_int(self, value):

if value > 255:

return False

if self.is_ascii() and value > 127:

return False

return (1 << (value & 0x1f)) & self.charbitmask()

def getslice(self, start, stop):

assert 0 <= start <= stop

return LiteralStringNode(self.s[start:stop])

def find_int(self, what, start, stop):

if not self.can_contain_int(what):

return -1

return self.s.find(chr(what), start, stop)

def _concat(self, other):

if (isinstance(other, LiteralStringNode) and

len(other.s) + len(self.s) < NEW_NODE_WHEN_LENGTH):

return LiteralStringNode(self.s + other.s)

elif (isinstance(other, LiteralUnicodeNode) and

len(other.u) + len(self.s) < NEW_NODE_WHEN_LENGTH and

len(self.s) < CONVERT_WHEN_SMALLER):

return LiteralUnicodeNode(self.s.decode("latin-1") + other.u)

return BinaryConcatNode(self, other)

def dot(self, seen, toplevel=False):

if self in seen:

return

seen[self] = True

addinfo = str(self.s).replace('"', "'") or "_"

if len(addinfo) > 10:

addinfo = addinfo[:3] + "..." + addinfo[-3:]

yield ('"%s" [shape=box,label="length: %s\\n%s"];' % (

id(self), len(self.s),

repr(addinfo).replace('"', '').replace("\\", "\\\\")))

LiteralStringNode.EMPTY = LiteralStringNode("")

LiteralStringNode.PREBUILT = [LiteralStringNode(chr(i)) for i in range(256)]

del i

class LiteralUnicodeNode(LiteralNode):

def __init__(self, u):

assert isinstance(u, unicode)

self.u = u

def length(self):

return len(self.u)

def flatten_unicode(self):

return self.u

def is_ascii(self):

return False # usually not

def is_bytestring(self):

return False

def hash_part(self):

return self.additional_info().hash

def compute_additional_info(self):

additional_info = GlobalRopeInfo()

charbitmask = 0

partial_hash = 0

for c in self.u:

ordc = ord(c)

charbitmask |= intmask(1 << (ordc & 0x1F))

partial_hash = (1000003*partial_hash) + ordc

partial_hash = intmask(partial_hash)

partial_hash |= HIGHEST_BIT_SET

additional_info.charbitmask = intmask(charbitmask)

additional_info.hash = partial_hash

self._additional_info = additional_info

return additional_info

def charbitmask(self):

return self.additional_info().charbitmask

def getunichar(self, index):

return self.u[index]

def getint(self, index):

return ord(self.u[index])

def getrope(self, index):

ch = ord(self.u[index])

if ch < 256:

return LiteralStringNode.PREBUILT[ch]

if len(self.u) == 1:

return self

return LiteralUnicodeNode(unichr(ch))

def can_contain_int(self, value):

return (1 << (value & 0x1f)) & self.charbitmask()

def getslice(self, start, stop):

assert 0 <= start <= stop

return LiteralUnicodeNode(self.u[start:stop])

def find_int(self, what, start, stop):

if not self.can_contain_int(what):

return -1

return self.u.find(unichr(what), start, stop)

def _concat(self, other):

if (isinstance(other, LiteralUnicodeNode) and

len(other.u) + len(self.u) < NEW_NODE_WHEN_LENGTH):

return LiteralUnicodeNode(self.u + other.u)

elif (isinstance(other, LiteralStringNode) and

len(other.s) + len(self.u) < NEW_NODE_WHEN_LENGTH and

len(other.s) < CONVERT_WHEN_SMALLER):

return LiteralUnicodeNode(self.u + other.s.decode("latin-1"))

return BinaryConcatNode(self, other)

def dot(self, seen, toplevel=False):

if self in seen:

return

seen[self] = True

addinfo = repr(self.u).replace('"', "'") or "_"

if len(addinfo) > 10:

addinfo = addinfo[:3] + "..." + addinfo[-3:]

yield ('"%s" [shape=box,label="length: %s\\n%s"];' % (

id(self), len(self.u),

repr(addinfo).replace('"', '').replace("\\", "\\\\")))

def make_binary_get(getter):

def get(self, index):

while isinstance(self, BinaryConcatNode):

llen = self.left.length()

if index >= llen:

self = self.right

index -= llen

else:

self = self.left

return getattr(self, getter)(index)

return get

class BinaryConcatNode(StringNode):

def __init__(self, left, right, balanced=False):

self.left = left

self.right = right

try:

self.len = ovfcheck(left.length() + right.length())

except OverflowError:

raise

self._depth = 0

# XXX the balance should become part of the depth

self.balanced = balanced

if balanced:

self.balance_known = True

else:

self.balance_known = False

def is_ascii(self):

return self.additional_info().is_ascii

def is_bytestring(self):

return self.additional_info().is_bytestring

def check_balanced(self):

if self.balance_known:

return self.balanced

# balance calculation

# XXX improve?

if not self.left.check_balanced() or not self.right.check_balanced():

balanced = False

else:

balanced = (find_fib_index(self.len // (NEW_NODE_WHEN_LENGTH / 2)) >=

self._depth)

self.balanced = balanced

self.balance_known = True

return balanced

def length(self):

return self.len

def depth(self):

depth = self._depth

if not depth:

depth = self._depth = max(self.left.depth(),

self.right.depth()) + 1

return depth

getchar = make_binary_get("getchar")

getunichar = make_binary_get("getunichar")

getint = make_binary_get("getint")

getrope = make_binary_get("getrope")

def can_contain_int(self, value):

if self.is_bytestring() and value > 255:

return False

if self.is_ascii() and value > 127:

return False

return (1 << (value & 0x1f)) & self.charbitmask()

def getslice(self, start, stop):

if start == 0:

if stop == self.length():

return self

return getslice_left(self, stop)

if stop == self.length():

return getslice_right(self, start)

return concatenate(

getslice_right(self.left, start),

getslice_left(self.right, stop - self.left.length()))

def flatten_string(self):

f = fringe(self)

return "".join([node.flatten_string() for node in f])

def flatten_unicode(self):

f = fringe(self)

return u"".join([node.flatten_unicode() for node in f])

def hash_part(self):

return self.additional_info().hash

def compute_additional_info(self):

leftaddinfo = self.left.additional_info()

rightaddinfo = self.right.additional_info()

additional_info = leftaddinfo.combine(rightaddinfo,

self.right.length())

self._additional_info = additional_info

return additional_info

def charbitmask(self):

return self.additional_info().charbitmask

def rebalance(self):

if self.balanced:

return self

return rebalance([self], self.len)

def _concat(self, other):

if isinstance(other, LiteralNode):

r = self.right

if isinstance(r, LiteralNode):

return BinaryConcatNode(self.left,

r._concat(other))

return BinaryConcatNode(self, other)

def dot(self, seen, toplevel=False):

if self in seen:

return

seen[self] = True

if toplevel:

addition = ", fillcolor=red"

elif self.check_balanced():

addition = ", fillcolor=yellow"

else:

addition = ""

yield '"%s" [shape=octagon,label="+\\ndepth=%s, length=%s"%s];' % (

id(self), self.depth(), self.len, addition)

for child in [self.left, self.right]:

yield '"%s" -> "%s";' % (id(self), id(child))

for line in child.dot(seen):

yield line

def concatenate(node1, node2):

if node1.length() == 0:

return node2

if node2.length() == 0:

return node1

result = node1._concat(node2)

if rebalance and result.depth() > MAX_DEPTH: #XXX better check

return result.rebalance()

return result

def getslice(node, start, stop, step, slicelength=-1):

if slicelength == -1:

# XXX for testing only

slicelength = len(xrange(start, stop, step))

start, stop, node = find_straddling(node, start, stop)

if step != 1:

iter = SeekableItemIterator(node)

iter.seekforward(start)

if node.is_bytestring():

result = [iter.nextchar()]

for i in range(slicelength - 1):

iter.seekforward(step - 1)

result.append(iter.nextchar())

return rope_from_charlist(result)

else:

result = [iter.nextunichar()]

for i in range(slicelength - 1):

iter.seekforward(step - 1)

result.append(iter.nextunichar())

return rope_from_unicharlist(result)

return node.getslice(start, stop)

def getslice_one(node, start, stop):

start, stop, node = find_straddling(node, start, stop)

return node.getslice(start, stop)

def find_straddling(node, start, stop):

while 1:

if isinstance(node, BinaryConcatNode):

llen = node.left.length()

if start >= llen:

node = node.right

start = start - llen

stop = stop - llen

continue

if stop <= llen:

node = node.left

continue

return start, stop, node

def getslice_right(node, start):

while 1:

if start == 0:

return node

if isinstance(node, BinaryConcatNode):

llen = node.left.length()

if start >= llen:

node = node.right

start = start - llen

continue

else:

return concatenate(getslice_right(node.left, start),

node.right)

return node.getslice(start, node.length())

def getslice_left(node, stop):

while 1:

if stop == node.length():

return node

if isinstance(node, BinaryConcatNode):

llen = node.left.length()

if stop <= llen:

node = node.left

continue

else:

return concatenate(node.left,

getslice_left(node.right, stop - llen))

return node.getslice(0, stop)

def multiply(node, times):

if times <= 0:

return LiteralStringNode.EMPTY

if times == 1:

return node

twopower = node

number = 1

result = None

while number <= times:

if number & times:

if result is None:

result = twopower

elif result.length() < CONCATENATE_WHEN_MULTIPLYING:

result = concatenate(result, twopower)

else:

result = BinaryConcatNode(result, twopower)

try:

number = ovfcheck(number * 2)

except OverflowError:

break

if twopower.length() < CONCATENATE_WHEN_MULTIPLYING:

twopower = concatenate(twopower, twopower)

else:

twopower = BinaryConcatNode(twopower, twopower)

return result

def join(node, l):

if node.length() == 0:

return rebalance(l)

nodelist = [None] * (2 * len(l) - 1)

length = 0

for i in range(len(l)):

nodelist[2 * i] = l[i]

length += l[i].length()

for i in range(len(l) - 1):

nodelist[2 * i + 1] = node

length += (len(l) - 1) * node.length()

return rebalance(nodelist, length)

def rebalance(nodelist, sizehint=-1):

if sizehint < 0:

sizehint = 0

for node in nodelist:

sizehint += node.length()

if sizehint == 0:

return LiteralStringNode.EMPTY

nodelist.reverse()

# this code is based on the Fibonacci identity:

# sum(fib(i) for i in range(n+1)) == fib(n+2)

l = [None] * (find_fib_index(sizehint) + 2)

stack = nodelist

empty_up_to = len(l)

a = b = sys.maxint

first_node = None

while stack:

curr = stack.pop()

while isinstance(curr, BinaryConcatNode) and not curr.check_balanced():

stack.append(curr.right)

curr = curr.left

currlen = curr.length()

if currlen == 0:

continue

if currlen < a:

# we can put 'curr' to its preferred location, which is in

# the known empty part at the beginning of 'l'

a, b = 1, 2

empty_up_to = 0

while not (currlen < b):

empty_up_to += 1

a, b = b, a+b

else:

# sweep all elements up to the preferred location for 'curr'

while not (currlen < b and l[empty_up_to] is None):

if l[empty_up_to] is not None:

curr = l[empty_up_to]._concat(curr)

l[empty_up_to] = None

currlen = curr.length()

else:

empty_up_to += 1

a, b = b, a+b

if empty_up_to == len(l):

return curr

l[empty_up_to] = curr

first_node = curr

# sweep all elements

curr = first_node

for index in range(empty_up_to + 1, len(l)):

if l[index] is not None:

curr = BinaryConcatNode(l[index], curr)

assert curr is not None

return curr

# __________________________________________________________________________

# construction from normal strings

def rope_from_charlist(charlist):

nodelist = []

size = 0

for i in range(0, len(charlist), NEW_NODE_WHEN_LENGTH):

chars = charlist[i: min(len(charlist), i + NEW_NODE_WHEN_LENGTH)]

nodelist.append(LiteralStringNode("".join(chars)))

size += len(chars)

return rebalance(nodelist, size)

def rope_from_unicharlist(charlist):

nodelist = []

length = len(charlist)

if not length:

return LiteralStringNode.EMPTY

i = 0

while i < length:

unichunk = []

while i < length:

c = ord(charlist[i])

if c < 256:

break

unichunk.append(unichr(c))

i += 1

if unichunk:

nodelist.append(LiteralUnicodeNode(u"".join(unichunk)))

strchunk = []

while i < length:

c = ord(charlist[i])

if c >= 256:

break

strchunk.append(chr(c))

i += 1

if strchunk:

nodelist.append(LiteralStringNode("".join(strchunk)))

return rebalance(nodelist, length)

def rope_from_unicode(uni):

nodelist = []

length = len(uni)

if not length:

return LiteralStringNode.EMPTY

i = 0

while i < length:

start = i

while i < length:

c = ord(uni[i])

if c < 256:

break

i += 1

if i != start:

nodelist.append(LiteralUnicodeNode(uni[start:i]))

start = i

while i < length:

c = ord(uni[i])

if c >= 256:

break

i += 1

if i != start:

nodelist.append(LiteralStringNode(uni[start:i].encode("latin-1")))

return rebalance(nodelist, length)

def rope_from_unichar(unichar):

intval = ord(unichar)

if intval > 256:

return LiteralUnicodeNode(unichar)

return LiteralStringNode.PREBUILT[intval]

# __________________________________________________________________________

# searching

def find_int(node, what, start=0, stop=-1):

offset = 0

length = node.length()

if stop == -1:

stop = length

if start != 0 or stop != length:

newstart, newstop, node = find_straddling(node, start, stop)

offset = start - newstart

start = newstart

stop = newstop

assert 0 <= start <= stop

if isinstance(node, LiteralNode):

pos = node.find_int(what, start, stop)

if pos == -1:

return pos

return pos + offset

if not node.can_contain_int(what):

return -1

# invariant: stack should only contain nodes that can contain the int what

stack = [node]

i = 0

while stack:

curr = stack.pop()

while isinstance(curr, BinaryConcatNode):

if curr.left.can_contain_int(what):

if curr.right.can_contain_int(what):

stack.append(curr.right)

curr = curr.left

else:

i += curr.left.length()

# if left cannot contain what, then right must contain it

curr = curr.right

nodelength = curr.length()

fringenode = curr

if i + nodelength <= start:

i += nodelength

continue

searchstart = max(0, start - i)

searchstop = min(stop - i, nodelength)

if searchstop <= 0:

return -1

assert isinstance(fringenode, LiteralNode)

pos = fringenode.find_int(what, searchstart, searchstop)

if pos != -1:

return pos + i + offset

i += nodelength

return -1

def find(node, subnode, start=0, stop=-1):

len1 = node.length()

len2 = subnode.length()

if stop > len1 or stop == -1:

stop = len1

if stop - start < 0:

return -1

if len2 == 1:

return find_int(node, subnode.getint(0), start, stop)

if len2 == 0:

return start

if len2 > stop - start:

return -1

restart = construct_restart_positions_node(subnode)

return _find_node(node, subnode, start, stop, restart)

def _find_node(node, subnode, start, stop, restart):

len2 = subnode.length()

m = start

iter = SeekableItemIterator(node)

iter.seekforward(start)

c = iter.nextint()

i = 0

subiter = SeekableItemIterator(subnode)

d = subiter.nextint()

while m + i < stop:

if c == d:

i += 1

if i == len2:

return m

d = subiter.nextint()

if m + i < stop:

c = iter.nextint()

else:

# mismatch, go back to the last possible starting pos

if i == 0:

m += 1

if m + i < stop:

c = iter.nextint()

else:

e = restart[i - 1]

new_m = m + i - e

assert new_m <= m + i

seek = m + i - new_m

if seek:

iter.seekback(m + i - new_m)

c = iter.nextint()

m = new_m

subiter.seekback(i - e + 1)

d = subiter.nextint()

i = e

return -1

def construct_restart_positions_node(node):

length = node.length()

restart = [0] * length

restart[0] = 0

i = 1

j = 0

iter1 = ItemIterator(node)

iter1.nextint()

c1 = iter1.nextint()

iter2 = SeekableItemIterator(node)

c2 = iter2.nextint()

while 1:

if c1 == c2:

j += 1

if j < length:

c2 = iter2.nextint()

restart[i] = j

i += 1

if i < length:

c1 = iter1.nextint()

else:

break

elif j>0:

new_j = restart[j-1]

assert new_j < j

iter2.seekback(j - new_j + 1)

c2 = iter2.nextint()

j = new_j

else:

restart[i] = 0

i += 1

if i < length:

c1 = iter1.nextint()

else:

break

j = 0

iter2 = SeekableItemIterator(node)

c2 = iter2.nextint()

return restart

def view(objs):

from dotviewer import graphclient

content = ["digraph G{"]

seen = {}

for i, obj in enumerate(objs):

if obj is None:

content.append(str(i) + ";")

else:

content.extend(obj.dot(seen, toplevel=True))

content.append("}")

p = py.test.ensuretemp("automaton").join("temp.dot")

p.write("\n".join(content))

graphclient.display_dot_file(str(p))

# __________________________________________________________________________

# iteration

class FringeIterator(object):

def __init__(self, node):

self.stack = [node]

def next(self):

while self.stack:

curr = self.stack.pop()

while 1:

if isinstance(curr, BinaryConcatNode):

self.stack.append(curr.right)

curr = curr.left

else:

return curr

raise StopIteration

def _seekforward(self, length):

"""seek forward up to n characters, returning the number remaining chars.

experimental api"""

curr = None

while self.stack:

curr = self.stack.pop()

if length < curr.length():

break

length -= curr.length()

else:

raise StopIteration

while isinstance(curr, BinaryConcatNode):

left_length = curr.left.length()

if length < left_length:

self.stack.append(curr.right)

curr = curr.left

else:

length -= left_length

curr = curr.right

self.stack.append(curr)

return length

def fringe(node):

result = []

iter = FringeIterator(node)

while 1:

try:

result.append(iter.next())

except StopIteration:

return result

class ReverseFringeIterator(object):

def __init__(self, node):

self.stack = [node]

def next(self):

while self.stack:

curr = self.stack.pop()

while 1:

if isinstance(curr, BinaryConcatNode):

self.stack.append(curr.left)

curr = curr.right

else:

return curr

raise StopIteration

class ItemIterator(object):

def __init__(self, node, start=0):

self.iter = FringeIterator(node)

self.node = None

self.nodelength = 0

self.index = 0

if start:

self._advance_to(start)

def _advance_to(self, index):

self.index = self.iter._seekforward(index)

self.node = self.iter.next()

self.nodelength = self.node.length()

def getnode(self):

node = self.node

if node is None:

while 1:

node = self.node = self.iter.next()

nodelength = self.nodelength = node.length()

if nodelength != 0:

self.index = 0

return node

return node

def advance_index(self):

index = self.index

if index == self.nodelength - 1:

self.node = None

else:

self.index = index + 1

def nextchar(self):

node = self.getnode()

result = node.getchar(self.index)

self.advance_index()

return result

def nextunichar(self):

node = self.getnode()

result = node.getunichar(self.index)

self.advance_index()

return result

def nextrope(self):

node = self.getnode()

result = node.getrope(self.index)

self.advance_index()

return result

def nextint(self):

node = self.getnode()

result = node.getint(self.index)

self.advance_index()

return result

class ReverseItemIterator(object):

def __init__(self, node):

self.iter = ReverseFringeIterator(node)

self.node = None