implementation of http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.46.7829 for plain lambda calculus

flow.py

#!/usr/bin/env python3
from util import *
from lam import *
from ty import *

def extract(t, *kind):
    if type(t) is Lam:
        res = extract(t.args[1], *kind)
        if Lam in kind: res.add(t)
        if Var in kind: res.add(t.args[0])
        return res
    elif type(t) is App:
        res = extract(t.args[0], *kind) | extract(t.args[1], *kind)
        if App in kind: res.add(t)
        return res
    else:
        res = set()
        if Var in kind: res.add(t)
        return res

@curry
def flow_iter(apps, flows):
    print('...')
    flows = multimap(transclos(flows))
    res = set()
    for a in apps:
        f,t = a.args
        for s in flows.get(f,()):
            if type(s) is Lam:
                (x,u) = s.args
                res.add((x,t))
                res.add((App(f,t),u))
    return res

def flow(term):
    apps = extract(term, App)
    lams = extract(term, Lam)
    x0 = {(t,t) for t in lams}
    flows = fixset(flow_iter(apps))(x0)
    flows = transclos(flows)
    mmap = multimap(flows)
    flows = {(k,v) for k,v in flows if type(v) is Lam and len(mmap[v]) == 1}
    return multimap(flows)

def scott(n):
    if n == 0:
        return Var('ZERO')
    else:
        return App('SUCC', scott(n-1))

def pair(x,y):
    return App(App('PAIR', x), y)

class UF:
    def __init__(self):
        self.find = {}
    def union(self, xs):
        #print('UNION', xs)
        u = xs.union(*[self.find[x] for x in xs if x in self.find])
        for e in u: self.find[e] = u
    def extract(self):
        return set(tuple(s) for s in self.find.values())

def main():
    #t = Lam(Var(0), App(Lam(Var(1), App(Var(1), Var(1))), Var(0)))
    
    t = let(('ZERO', ['z0','s0', 'z0']),
            ('SUCC', ['n','zs','ss', ('ss', 'n')]),
            ('OMEGA', ['o', ('o','o')]),
            ('Y', ['g', ('OMEGA', ['y', ('g', ('y','y'))])]),
            ('ADD', ('Y', ['add','a','b',
                        ('a', 'b', ['m','za','sa',
                            ('sa', ('add', 'm', 'b'))])])),
            ('TRUE', ['xt','yt', 'xt']),
            ('FALSE', ['xf','yf', 'yf']),
            #('PRN', ['num', ('num', 'FALSE',
            #    ['num1','p', ('p', 'TRUE', ('PRN', 'num1'))])]),
        ('ADD', scott(10), ('ADD', scott(15), scott(0))))
        #('PRN', ('ADD', scott(3), ('ADD', scott(4), scott(0)))))
    '''
    t = let(('ZERO', Lam('z0', Lam('s0', 'z0'))),
            ('SUCC', Lam('n', Lam('zs', Lam('ss', App('ss', 'n'))))),
            ('OMEGA', Lam('o', App('o', 'o'))),
            ('Y', Lam('g', App('OMEGA', Lam('y', App('g', App('y','y')))))),
            ('TRUE', Lam('xt', Lam('_t', 'xt'))),
            ('FALSE', Lam('_f', Lam('yf', 'yf'))),
            ('PAIR', Lam('e1', Lam('e2', Lam('p',
                App(App('p', 'e1'), 'e2'))))),
            ('FST', Lam('p1', App('p1', 'TRUE'))),
            ('SND', Lam('p2', App('p2', 'FALSE'))),
            ('SUC', Lam('t', Lam('f', Lam('xs', pair(
                App('FST', 'x'), App(App('t', 'f'), App('SND', 'xs'))))))),
            ('SIEVE', App('Y', Lam('sieve', Lam('n',
                pair('TRUE', App(App('Y', App('SUC', 'n')),
                    App('SIEVE', App('SUC', 'n')))))))),
            ('PRIMES', pair('FALSE', pair('FALSE', App('SIEVE',
                Lam('cont', Lam('zs', pair('FALSE',
                    App('cont', App('SND', 'ys'))))))))),
            ('INDEX', App('Y', Lam('index', Lam('seq', Lam('i',
                App(App('i',
                    App('FST', 'seq')),
                    Lam('j', App('index', App(App('SND', 'seq'), 'j'))))))))),
        App(App('INDEX', 'PRIMES'), scott(10)))
    
    t = let(('OMEGA', Lam('o', App('o', 'o'))),
            ('Y', Lam('g', App('OMEGA', Lam('y', App('g', App('y','y')))))),
            ('TRUE', Lam('xt', Lam('_t', 'xt'))),
            ('FALSE', Lam('_f', Lam('yf', 'yf'))),
            ('PAIR', Lam('e1', Lam('e2', Lam('p',
                App(App('p', 'e1'), 'e2'))))),
            ('FST', Lam('p1', App('p1', 'TRUE'))),
            ('SND', Lam('p2', App('p2', 'FALSE'))),
            ('SUC', Lam('t', Lam('f', Lam('xs', pair(
                App('FST', 'x'), App(App('t', 'f'), App('SND', 'xs'))))))),
            ('SIEVE', Lam('n',
                pair('TRUE', App(App('Y', App('SUC', 'n')),
                    App('SIEVE', App('SUC', 'n')))))),
            ('PRIMES', pair('FALSE', pair('FALSE', App('SIEVE',
                Lam('cont', Lam('zs', pair('FALSE',
                    App('cont', App('SND', 'ys'))))))))),
        Var('PRIMES'))
    '''

    #print(t); print()

    flows = flow(t)

    # extract types from flow graph
    tyd = {}
    tyof = {}
    for term,vs in flows.items():
        if len(vs) == 1 and term in vs: continue
        #if len(str(term)) > 15: continue
        tys = ''
        for i,v in enumerate(sorted(vs)):
            v = v.subst(term, Var('…'))
            v = v.debruijn()
            s = str(v)
            if len(vs) == 1: # potential inlining
                tys = ":= %s" % v
            elif not i:
                tys = "= { %s" % s[1:-1]
            else:
                tys += " | %s" % s[1:-1]
        if len(vs) > 1: tys += ' }'
        #print(term, '\t', tys)
        if tys not in tyd: tyd[tys] = set()
        tyd[tys].add(str(term))
        #print(term, tys)
        if str(term) not in tyof or '…' not in tyof[str(term)]:
            tyof[str(term)] = tys
        tys1 = tys.replace('…',str(term))
        if tys1 not in tyd: tyd[tys1] = set()
        tyd[tys1].add(str(term))

    # construct type-equivalence classes over terms
    uf = UF()
    for tys,terms in tyd.items():
        #if len(terms) < 2: continue
        uf.union(terms)
        '''
        ts = list(terms)
        print('UNION', ts)
        ts.sort()
        ts.sort(key=len)
        t0 = ts[0]
        uf.make(t0)
        for t in ts[1:]:
            uf.make(t)
            uf.union(t,t0)
        '''
    for ts in uf.extract():
        #print('*', t0)
        print('---')
        ty = ''
        ts = sorted(list(ts), key=len)
        for t in ts:
            print(t)
            if '…' not in ty: ty = tyof[t]
        print('\t', ty)
        print()
        #print('\t===', terms[0])
        #print('\t', tys)

    '''
    print()
    #n = App('add', 'm', 'b')
    #for v in (Var('add'),):
    for v in extract(t, Var):
        #v = Var('s0')
        ty = infer(flows,v)
        #print(v)
        #if type(ty) is TyRec: print_('new')
        name = slugify(str(v), separator='_')
        print_('data T%s = ' % name)
        pprint(ty)
    '''
if __name__ == '__main__': main()

flow.txt

let	 ZERO 	= (λz0 s0. z0)
	 SUCC 	= (λn zs ss. ss n)
	 OMEGA 	= (λo. o o)
	 Y 	= (λg. OMEGA (λy. g (y y)))
	 ADD 	= (Y (λadd a b. a b (λm za sa. sa (add m b))))
	 TRUE 	= (λxt yt. xt)
	 FALSE 	= (λxf yf. yf)
in (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

...
...
...
...
...
...
...
...
...
...
...
...
...
---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n))
	 := (λλλλλ.2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

---
(ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))))
(ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))))
	 := (λ.a 0̲ (λλλ.0̲ (add 2̲ 3̲)))

---
Y
	 := (λ.OMEGA (λ.1̲ (0̲ 0̲)))

---
ss
s0
	 := (λλλ.0̲ (add 2̲ b))

---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n) (λo. o o))
	 := (λλλλ.2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n) (λo. o o) (λg. OMEGA (λy. g (y y))) (Y (λadd a b. a b (λm za sa. sa (add m b)))))
	 := (λλ.ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

---
(a b)
	 = { λ.z0 | λ.0̲ n }

---
y
o
	 := (λ.g (0̲ 0̲))

---
add
ADD
(y y)
(o o)
(g (y y))
(OMEGA (λy. g (y y)))
(Y (λadd a b. a b (λm za sa. sa (add m b))))
	 := (λλ.1̲ 0̲ (λλλ.0̲ (… 2̲ 3̲)))

---
ZERO
TRUE
	 := (λλ.1̲)

---
FALSE
	 := (λλ.0̲)

---
(SUCC ZERO)
(SUCC (SUCC ZERO))
(SUCC (SUCC (SUCC ZERO)))
(SUCC (SUCC (SUCC (SUCC ZERO))))
(SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))
(SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))))))))
	 := (λλ.0̲ n)

---
SUCC
	 := (λλλ.0̲ 2̲)

---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n) (λo. o o) (λg. OMEGA (λy. g (y y))) (Y (λadd a b. a b (λm za sa. sa (add m b)))) (λxt yt. xt))
	 := (λ.ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

---
OMEGA
	 := (λ.0̲ 0̲)

---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0))
	 := (λλλλλλ.2̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ ZERO)))))))))) (2̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ (5̲ ZERO))))))))))))))) ZERO))

---
g
	 := (λλλ.1̲ 0̲ (λλλ.0̲ (5̲ 2̲ 3̲)))

---
n
a
b
m
zs
z0
(add m b)
(a b (λm za sa. sa (add m b)))
(ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)
(ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n) (λo. o o) (λg. OMEGA (λy. g (y y))) (Y (λadd a b. a b (λm za sa. sa (add m b)))) (λxt yt. xt) (λxf yf. yf))
	 = { λλ.1̲ | λλ.0̲ … }

---
(add m)
	 := (λ.a 0̲ (λλλ.0̲ (… 3̲)))

---
((λZERO SUCC OMEGA Y ADD TRUE FALSE. ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (ADD (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO)) (λz0 s0. z0) (λn zs ss. ss n) (λo. o o) (λg. OMEGA (λy. g (y y))))
	 := (λλλ.2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO)))))))))) (2̲ (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC (SUCC ZERO))))))))))))))) ZERO))

---
(ss n)
	 := (λλ.0̲ (add m b))

lam.py

import sys
from util import *
from functools import total_ordering

@total_ordering
class Term:
    def __init__(self, *args):
        self.args = tuple(Var(a) if type(a) in (str,int) else a for a in args)
    def __repr__(self):
        return str(self)
    def __eq__(self, other):
        return type(self) is type(other) and self.args == other.args
    def __lt__(self, other):
        s0 = str(self)
        s1 = str(other)
        return len(s0) < len(s1) or s0 < s1
    def __hash__(self):
        return hash(self.args)
class Var(Term):
    def __init__(self, x):
        self.x = x
        self.args = (self.x,)
    def subst(self, this, that):
        if self == this: return that
        return self
    def debruijn(self, depth=0, d=None):
        if d is None: d = {}
        if self in d: return depth - d[self] - 1
        return self
    def __str__(self):
        if type(self.x) is int:
            return str(self.x) + '\u0332'
        return self.x
class Lam(Term):
    def __init__(self, *args):
        if len(args) > 2:
            args = (args[0], Lam(*args[1:]))
        Term.__init__(self, *args)
    def subst(self, this, that):
        if self == this: return that
        if self.args[0] == this: return self
        return Lam(self.args[0].subst(this,that),
                   self.args[1].subst(this,that))
    def debruijn(self, depth=0, d=None):
        if d is None: d = {}
        else: d = dict(d.items())
        d[self.args[0]] = depth
        return Lam('', self.args[1].debruijn(depth+1,d))
    def __str__(self):
        if self.args[0].x:
            s0 = str(self.args[0])
            s1 = str(self.args[1])
            if type(self.args[1]) is Lam:
                s0 = s0 + ' '
                s1 = s1[2:-1]
            else:
                s0 = s0 + '. '
            if type(self.args[1]) is App:
                s1 = s1[1:-1]
            return "(λ%s%s)" % (s0,s1)
        else:
            s = str(self.args[1])
            if type(self.args[1]) is not Var: s = s[1:-1]
            if type(self.args[1]) is not Lam: s = '.' + s
                
            return "(λ%s)" % s
class App(Term):
    def __init__(self, *args):
        if len(args) > 2:
            args = (App(*args[:-1]), args[-1])
        Term.__init__(self, *args)
    def subst(self, this, that):
        if self == this: return that
        return App(self.args[0].subst(this,that),
                   self.args[1].subst(this,that))
    def debruijn(self, depth=0, d=None):
        if d is None: d = {}
        return App(self.args[0].debruijn(depth,d),
                   self.args[1].debruijn(depth,d))
    def __str__(self):
        s0 = str(self.args[0])
        s1 = str(self.args[1])
        if type(self.args[0]) is App: s0 = s0[1:-1]
        #if type(self.args[0]) is not Var: s1 = "(%s)" % s1
        return "(%s %s)" % (s0,s1)

def expand(t):
    if type(t) is tuple:
        return App(*map(expand,t))
    elif type(t) is list:
        return Lam(*map(expand,t))
    else:
        return t
def let(*args, quiet=False):
    if not quiet:
        term = expand(args[-1])
        defs = args[:-1]
        print_('let', file=sys.stderr)
        for k,v in defs:
            print('\t', k, '\t=', expand(v), file=sys.stderr)
        print('in', term, file=sys.stderr)
        print(file=sys.stderr)

    #if len(args) == 1: return args[0]
    #return App(Lam(args[0][0], let(*args[1:], quiet=True)), expand(args[0][1]))
    lam = [p[0] for p in defs] + [term]
    app = [expand(p[1]) for p in defs]
    return App(Lam(*lam), *app)

ty.py

from util import *
from termcolor import colored
from slugify import slugify

class Type:
    def __init__(self, *args):
        self.args = tuple(args)
        self.hash = hash(self.args)
        self._str = None
    '''
    def __str__(self):
        return str(self.args)
    def __str__(self):
        if not self._str:
            self._str = '!!!RECURSION ERROR!!!'
            self._str = self.tostr()
        return self._str
    def __repr__(self):
        return str(self)
    '''
    def __eq__(self, other):
        return type(self) is type(other) and self.args == other.args
    def __hash__(self):
        return self.hash
    def __contains__(self, other):
        if self == other: return True
        for arg in self.args:
            if isinstance(arg,Type) and other in arg: return True
class TyVar(Type):
    def _repr_pretty_(self, p, cycle):
        #p.text("α" + colored("[%s]" % self.args[0], 'yellow'))
        p.text('T' + slugify(str(self.args[0]), separator='_'))
class TyFun(Type):
    def __init__(self, *args):
        Type.__init__(self, *args)
    def _repr_pretty_(self, p, cycle):
        p.text('(')
        p.pretty(self.args[0])
        p.text(" → ")
        p.pretty(self.args[1])
        p.text(')')
    def __str__(self):
        return "(%s → %s)" % self.args
class TyRec(Type):
    def _repr_pretty_(self, p, cycle):
        p.text("(μ")
        p.pretty(self.args[0])
        p.text('.')
        p.pretty(self.args[1])
        p.text(')')
class TySum(Type):
    def _repr_pretty_(self, p, cycle):
        with p.group(4,):
            if len(self.args) == 2:
                p.text('(Either')
            else:
                p.text("(S%d" % len(self.args))
            for i,ty in enumerate(self.args):
                p.breakable()
                p.text('(')
                p.pretty(ty)
                p.text(')')
            p.text(')')
class TyWild(Type):
    def _repr_pretty_(self, p, cycle):
        #p.text('*')
        #p.text('w%d' % (id(self) % 10000))
        p.text('*' + colored("[%s]" % self.args[0], 'blue'))
    def __str__(self):
        return "*[%s]" % self.args

def infer(flows, term, rec=(), depth=1, maxdepth=10):
    #if depth == maxdepth:
    #    print('! '*depth, term, 'STOP!')
    #   return TyVar(term)
    if term not in flows: return TyWild(term)
    #print('. '*depth, colored(term, 'green'))
    rec = {term}.union(rec)
    types = set()
    for lam in flows[term]:
        #print('lam =', lam)
        x,v = lam.args
        #print('x =', x)
        #print('v =', v)
        #if x in rec:
        #    #print('rec', x)
        #    ax = TyVar(x)
        #else:
        #    ax = infer(flows, x, rec, depth+1, maxdepth)
        ax = TyVar(x)
        if x not in flows: ax = TyWild(x)
        if v in rec:
            #print('rec', v)
            av = TyVar(v)
        else:
            av = infer(flows, v, rec, depth+1, maxdepth)
        if type(av) is TyRec: av = av.args[0]
        #if ax is not None and av is not None:
        t = TyFun(ax,av)
        #if type(ax) is TyWild and type(av) is TyWild: t = TyWild(TyFun(ax.args[0],av.args[0]))
        types.add(t)
        #print_('  '*depth, term, ': '); pprint(t)
    if len(types) == 0: return None
    if len(types) == 1:
        res = list(types)[0]
    else:
        res = TySum(*types)
    mu = TyVar(term)
    if mu in res:
        return TyRec(mu, res)
    else:
        return res

util.py

from collections import defaultdict
from functools import partial, wraps
from IPython.lib.pretty import pprint

def print_(*args,**kwargs):
    print(*args, end='', flush=True, **kwargs)

def curry(f):
    @wraps(f)
    def g(*args, **kwargs):
        return partial(f, *args, **kwargs)
    return g

@curry
def fix(f,x0):
    x = None
    while x0 != x:
        x = x0
        #print(x)
        x0 = f(x)
    return x

def fixset(f):
    return fix(lambda x: x.union(f(x)))

transclos = fixset(lambda x: {(a,d) for a,b in x for c,d in x if b == c})

def multimap(edges):
    d = defaultdict(set)
    for (u,v) in edges:
        d[u].add(v)
    return dict(d)