Gradient-friendly stack test

stack.py

from collections import namedtuple
import unittest

Compose = namedtuple("Compose", ["a", "b"])


class Stack(object):

    def __init__(self, compose_fn, n):
        # self.data[i] stores the top element of the stack after `i` updates
        self.data = [None for _ in range(2 * n)]

        # Queue of stack positions which should be used in merge operations
        self._merge_ptrs = [None for _ in range(n)]

        # Push pointer. Position at which next element should be inserted.
        self.i = 0

        # Merge queue pointer.
        self.j = -1

        self._c = compose_fn

    def push(self, x):
        self.data[self.i] = x

        # Add element to merge queue.
        self.j += 1
        self._merge_ptrs[self.j] = self.i

        # Advance the push pointer.
        self.i += 1

    def merge(self):
        # Elements to compose: 1) top of stack, 2) most recent item on merge
        # queue
        c1 = self.data[self.i - 1]
        c2 = self.data[self._merge_ptrs[self.j - 1]]
        self.data[self.i] = self._c(c1, c2)

        # Add composed element to the merge queue.
        self.j -= 1
        self._merge_ptrs[self.j] = self.i

        # Advance push pointer.
        self.i += 1

    def peek(self):
        final = filter(None, self.data)[-1]
        return final


class StackTest(unittest.TestCase):

    def test_pushes(self):
        stack = Stack(Compose, 5)

        stack.push("a")
        stack.push("b")
        stack.push("c")
        stack.push("d")

        self.assertEqual(stack.peek(), "d")

    def test_left_branching(self):
        stack = Stack(Compose, 5)

        stack.push("a")
        stack.push("b")
        stack.merge()
        stack.push("c")
        stack.merge()
        stack.push("d")
        stack.merge()

        self.assertEqual(stack.peek(),
                Compose("d", Compose("c", Compose("b", "a"))))

    def test_right_branching(self):
        stack = Stack(Compose, 5)

        stack.push("a")
        stack.push("b")
        stack.push("c")
        stack.push("d")
        stack.merge()
        stack.merge()
        stack.merge()

        self.assertEqual(stack.peek(),
                Compose(Compose(Compose("d", "c"), "b"), "a"))

    def test_stack(self):
        stack = Stack(Compose, 5)

        stack.push("a")
        stack.push("b")
        stack.push("c")
        stack.merge()
        stack.merge()
        stack.push("d")
        stack.merge()

        self.assertEqual(stack.peek(),
                Compose("d", Compose(Compose("c", "b"), "a")))

    def test_stack2(self):
        stack = Stack(Compose, 5)

        stack.push("a")
        stack.push("b")
        stack.merge()
        stack.push("c")
        stack.push("d")
        stack.merge()
        stack.merge()

        self.assertEqual(stack.peek(),
                Compose(Compose("d", "c"), Compose("b", "a")))


if __name__ == "__main__":
    unittest.main()