Irakli Safareli

Stack safe Function composition


Executing composition of many functions is not stack safe. For example:

// if we have a list of many functions
const fs = Array(100000).fill().map((_,idx) => x => x + idx)

// and we want to compose them
const compose = (f,g) => x => f(g(x))

// we can either use compose to create a new function which is
// composition of other functions, which is not stack safe
fs.reduce(compose)(1)

// or we can "interpret" the list of functions step by step 
// in a stack safe way
fs.reduceRight((arg,f) => f(arg),1) // 4999950001

To fix this issue we need a way to compose functions such that instead of creating another function (which when used requires extra stack frame leading to stack overflow), we collect those small functions as list of functions, which then could be executed in a stack safe way.

Here we create a special type Func a b which expresses computations from values of type a to values of type b (similar to (->)).

// type Func a b where
//   One  :: (a -> b) -> Func a b
//   Pipe :: Func a z -> Func z b -> Func a b

const Func = taggedSum({
  One: ['f'],
  Pipe: ['x', 'y'],
})

The representation we have is like a binary tree where data is on leafs. we could have used Linked List like structure instead, but Lists have bad performance during left associated concatenation, so using it would lead us to bad performance during composition. actually the bed performance issue could be fixed using Difference list (DList), but it represents elements of List using Functions and concatenation using Function composition, which could lead us to same stack overflow issue we want to solve. By representation Func with binary tree complexity of compose is O(1) and for lower it’s O(n) where n is number of functions in whole composition.

Now we can make it an instance of Category (actually we don’t yet have Category defined in Fantasy land but hope we’ll get it soon)

// :: (a -> b) -> Func a b
Func.lift = Func.One

// :: Ɐ a. Func a a
Func.id = Func.lift(a => a)

// :: Func b c ~> Func a b -> Func a c
Func.prototype.compose = function(f) {
  return Func.Pipe(f, this)
}

Now we can lift normal functions from a to b into Func a b and compose them! the composition builds up a tree of functions which could be inspected/interpreted in a stack safe way.

// here we take a Func value and then visit each node in stack safe
// way and build Array of Functions in correct order.
// complexity of this function is `O(1)`
//
// :: Func a z -> [a->b, b->c, ..., x->y, y->z]
const toArrayOfFunctions = function(root) {
  const fs = []
  const stack = []
  while (true) {
    if (root.tag == 'One'){
      fs.push(root.f)
      if (stack.length == 0) {
        break
      } else {
        root = stack.pop()
      }
    } else {
      stack.push(root.y)
      root = root.x
    }
  }
  return fs
}

// :: Func a b ~> () -> (a -> b)
Func.prototype.lower = function() {
  // we precompute `fs` here so that executing result multiple times
  // does not recomputes it again
  const fs = toArrayOfFunctions(this)
  return x => fs.reduce((arg, f) => f(arg), x)
}

// :: Func a b ~> () -> a -> b
Func.prototype.call = function(x) {
  return this.lower()(x)
}

To get feeling how it works lets see a small example:

const ab = Func.lift(a => [a, 'b'])
const bc = Func.lift(b => [b, 'c'])
const ac = bc.compose(ab)

const cd = Func.lift(c => [c, 'd'])
const de = Func.lift(d => [d, 'e'])
const ce = de.compose(cd)
const ae = ce.compose(ac)

const ez = Func.lift(e => [e, 'z'])
const az = ez.compose(ae)


console.log(toString(az)) /*
Func.Pipe(
  Func.Pipe(
    Func.Pipe(
      Func.One(a => [a, 'b']),
      Func.One(b => [b, 'c'])
    ),
    Func.Pipe(
      Func.One(c => [c, 'd']),
      Func.One(d => [d, 'e'])
    )
  ),
  Func.One(e => [e, 'z'])
) */

console.log(toString(toArrayOfFunctions(az))) /*
[ a => [a, 'b']
, b => [b, 'c']
, c => [c, 'd']
, d => [d, 'e']
, e => [e, 'z']
] */

console.log(toString(az.call('a'))) /*
[[[[["a", "b"], "c"], "d"], "e"], "z"] */

Now we can go back to our initial example and fix it:

const fs = Array(100000).fill()
  .map((_, idx) => Func.lift(x => x + idx))

const f = fs.reduce((ab, bc) => bc.compose(ab))

f.call(1)) // 4999950001

I think Func could be abstracted into a Free Category, which I tried to find but couldn’t, so there are things to explore in this direction. In general Free structures represent computations as data which then could be inspected and interpreted in a clever way. Func is a nice example of how this technique could be useful to fix stack size issue we had initially.

P.S. I came to this idea from this discussion

P.S.S. You might wanna take a look at safareli/purescript-stacksafe-function