jamcue.py

#!/bin/env python

import typing
from typing import Union, List, NamedTuple, Optional, Dict, Tuple, Iterable

atom = int
bloq = atom
patU = atom
patS = int   # signed integer
cell = tuple # Tuple[atom]  # tuple is used to model an immutable list. just list doesn't work b/c it is not `hash`able
noun = Union[atom, cell]
unit = Optional[noun]

def cue(a: atom) -> noun:
	"""
	++  cue
	  ~/  %cue
	  |=  a=@
	  ^-  *
	  =+  b=0
	  =+  m=`(map @ *)`~
	  =<  q
	  |-  ^-  [p=@ q=* r=(map @ *)]
	  ?:  =(0 (cut 0 [b 1] a))
		=+  c=(rub +(b) a)
		[+(p.c) q.c (~(put by m) b q.c)]
	  =+  c=(add 2 b)
	  ?:  =(0 (cut 0 [+(b) 1] a))
		=+  u=$(b c)
		=+  v=$(b (add p.u c), m r.u)
		=+  w=[q.u q.v]
		[(add 2 (add p.u p.v)) w (~(put by r.v) b w)]
	  =+  d=(rub c a)
	  [(add 2 p.d) (need (~(get by m) q.d)) m]
	"""
	a = normalize_noun(a)

	b = 0
	m: Dict[atom, noun] = {}
	TrapOutput = NamedTuple("TrapOutput", p=atom, q=noun, r=Dict[atom, noun])
	def cue_trap(a, b, m) -> TrapOutput:
	
		if 0 == cut(0, [b, 1], a):
			c = rub(b + 1, a)
			return TrapOutput(c.p + 1, c.q, {**m, b: c.q})  # update dict m immutably with new key b

		c = 2 + b
		if 0 == cut(0, [b + 1, 1], a):

			u = cue_trap(a, b=c, m=m)
			v = cue_trap(a, b=(u.p + c), m=u.r)
			w = [u.q, v.q]

			return TrapOutput((2 + (u.p + v.p)), w, {**v.r, b: w})  # dict unpacking

		d = rub(c, a)
		return TrapOutput((2 + d.p), need(m[d.q]), m)

	return normalize_noun(cue_trap(a, b, m).q)



def need(a: unit) -> noun:
	assert a is not None
	return a

# jam/cue actually works :P
def jam(a: noun) -> atom:
	"""
	++  jam
	  ~/  %jam
	  |=  a=*
	  ^-  @
	  =+  b=0
	  =+  m=`(map * @)`~
	  =<  q
	  |-  ^-  [p=@ q=@ r=(map * @)]
	  =+  c=(~(get by m) a)
	  ?~  c
		=>  .(m (~(put by m) a b))
		?:  ?=(@ a)
		  =+  d=(mat a)
		  [(add 1 p.d) (lsh 0 1 q.d) m]
		=>  .(b (add 2 b))
		=+  d=$(a -.a)
		=+  e=$(a +.a, b (add b p.d), m r.d)
		[(add 2 (add p.d p.e)) (mix 1 (lsh 0 2 (cat 0 q.d q.e))) r.e]
	  ?:  ?&(?=(@ a) (lte (met 0 a) (met 0 u.c)))
		=+  d=(mat a)
		[(add 1 p.d) (lsh 0 1 q.d) m]
	  =+  d=(mat u.c)
	  [(add 2 p.d) (mix 3 (lsh 0 2 q.d)) m]
	"""
	assert is_noun(a)

	b = 0
	m: Dict[noun, atom] = {}

	jam_counter = 0           # debug
	
	TrapOutput = NamedTuple("TrapOutput", p=atom, q=atom, r=Dict[noun, atom])
	def jam_trap(a, b, m) -> TrapOutput:
		nonlocal jam_counter  # debug
		jam_counter += 1      # debug
		a = normalize_noun(a) # needed to ensure that (2,) becomes 2. otherwise the impl breaks

		c = m.get(a)

		print(f"[{jam_counter}] a={a} b={b} m={m}")

		if c is None:		
			m[a] = b
			if type(a) == atom:
				print("c is none; a is atom")
				d = mat(a)
				return TrapOutput((1 + d.p), lsh(0, 1, d.q), m)
			else:
				print("c is none; a is cell")
				b = 2 + b
				d=jam_trap(a=cell_head(a), b=b, m=m)  # a = -.a (i.e. head of a)
				e=jam_trap(a=cell_tail(a), b=(b + d.p), m=d.r)
				return TrapOutput((2 + (d.p + e.p)), mix(1, lsh(0, 2, cat(0, d.q, e.q))), e.r)
		# c is not none, so it has to be a noun.
		elif type(a) == atom and (met(0, a) <= met(0, c)):
			print("a is atom; met(0, a) <= met(0, c)")
			d = mat(a)
			return TrapOutput(1 + d.p, lsh(0, 1, d.q), m)
		else:
			print("either a is a cell or met(0, a) > met(0, c)")
			d = mat(c)
			return TrapOutput(2 + d.p, mix(3, lsh(0, 2, d.q)), m)

	return jam_trap(a, b, m).q


#### custom utilities
def cell_head(a: cell) -> cell:
	if len(a) > 0: 
		return a[0]
	else:
		return ()

def cell_tail(a: cell):
	return a[1:]

def is_noun(a):
	# we'd need typing_inspect.is_generic_type(noun)
	# return isinstance(a, noun.__args__)
	return is_atom(a) or is_cell(a)

def is_atom(a):
	return isinstance(a, atom)

def is_cell(a):
	return isinstance(a, cell)


def normalize_noun(a: noun) -> noun:
	if type(a) == atom:
		return a
	elif len(a) == 0:
		# empty list is the same as zero in hoon
		return 0
	elif len(a) == 1:
		# one-element list is the same as the element in hoon
		return a[0]
	else:
		# recursively normalize all elements of list
		# return the noun converted to a tuple
		# with all of its elements normalized as well
		# this has the effect of converting all nested iterables into tuples
		"""
		In [27]: normalize_noun([1,2, [3,4], (5,)])                                                                           
		Out[27]: (1, 2, (3, 4), 5)
		"""
		return tuple(map(lambda x: normalize_noun(x), a))

def to_urbit_num(a: atom) -> str:
	s = []
	for i, dgt in enumerate(reversed(str(a))):
		if i != 0 and i % 3 == 0:
			s.append(".")
		s.append(dgt)

	return "".join(reversed(s))

def from_urbit_num(n: str) -> int:
	return int(n.replace(".", ""))

####################### jam/cue helpers
RubOutput = NamedTuple("RubOutput", p=atom, q=atom)
def rub(a: atom, b: atom) -> RubOutput:
	"""
	++  rub
	  ~/  %rub
	  |=  [a=@ b=@]
	  ^-  [p=@ q=@]
	  =+  ^=  c
		  =+  [c=0 m=(met 0 b)]
		  |-  ?<  (gth c m)
		  ?.  =(0 (cut 0 [(add a c) 1] b))
			c
		  $(c +(c))

	  ?:  =(0 c)
		[1 0]
	  =+  d=(add a +(c))
	  =+  e=(add (bex (dec c)) (cut 0 [d (dec c)] b))
	  [(add (add c c) e) (cut 0 [(add d (dec c)) e] b)]
	"""

	# pre ^= c
	c = 0
	m = met(0, b)
	def c_trap(a, b, c, m):
		assert not (c > m)
		if 0 == cut(0, [a + c, 1], b):
			return c_trap(a, b, c=c+1, m=m)
		else:
			return c
	
	c = c_trap(a, b, c, m)  # ^= c
	# post ^= c

	if 0 == c:
		return RubOutput(1, 0)

	d = a + (c + 1)
	e = bex(c - 1) + cut(0, [d, c - 1], b)
	return RubOutput((c + c) + e, cut(0, [d + (c - 1), e], b))



def cut(a: bloq, bc: list, d: atom):
	"""
	++  cut
	  ~/  %cut
	  |=  [a=bloq [b=@u c=@u] d=@]
	  (end a c (rsh a b d))
	"""
	b, c = bc
	return end(a, c, rsh(a, b, d))




MatReturn = NamedTuple("MatReturn", p=atom, q=atom)
def mat(a: atom) -> MatReturn:
	"""
	Length-encode

	Produces a cell whose tail q is atom a with a bit representation of its length prepended to it (as the least significant bits). The head p is the length of q in bits.
	"""
	"""
	++  mat
		~/  %mat
		|=  a=@
		^-  [p=@ q=@]
		?:  =(0 a)
			[1 1]
		=+  b=(met 0 a)
		=+  c=(met 0 b)
		:-  (add (add c c) b)
		(cat 0 (bex c) (mix (end 0 (dec c) b) (lsh 0 (dec c) a)))
	"""

	if a == 0:
		return MatReturn(1, 1)
	
	b = met(0, a)
	c = met(0, b)

	return MatReturn(
		p=(c+c) + b,
		q=cat( 
			0,
			bex(c),
			mix(
				end(0, (c - 1), b),
				lsh(0, (c - 1), a)
			)
		)
	)


def met(a: bloq, b: atom) -> atom:
	"""
	++  met
	  ~/  %met
	  |=  [a=bloq b=@]
	  ^-  @
	  =+  c=0
	  |-
	  ?:  =(0 b)  c
	  $(b (rsh a 1 b), c +(c))
	"""
	c = 0
	def met_trap(a, b, c):
		if 0 == b: return c
		return met_trap(a, b=rsh(a, 1, b), c=c+1)

	return met_trap(a, b, c)


def cat(a: bloq, b: atom, c: atom):
	"""
	++  cat
	  ~/  %cat
	  |=  [a=bloq b=@ c=@]
	  (add (lsh a (met a b) c) b)
	"""
	return lsh(a, met(a, b), c) + b



def mix(a: atom, b: atom) -> atom:
	"""
	Binary XOR

	Produces the bitwise logical XOR of two atoms, a and b, producing an atom.
	"""

	"""
	++  mix
	  ~/  %mix
	  |=  [a=@ b=@]
	  ^-  @
	  =+  [c=0 d=0]
	  |-
	  ?:  ?&(=(0 a) =(0 b))  d
	  %=  $
		a   (rsh 0 1 a)
		b   (rsh 0 1 b)
		c   +(c)
		d   (add d (lsh 0 c =((end 0 1 a) (end 0 1 b))))
	  ==

	"""
	c = 0
	d = 0
	def mix_trap(a, b, c, d):
		if (0 == a and 0 == b): return d
		return mix_trap(
			a=rsh(0, 1, a),
			b=rsh(0, 1, b),
			c=(c+1),
			d=(
				d + lsh(
						0,
						c,
						int(loob2bool(end(0, 1, a) == end(0, 1, b)))
					)
			)
		)

	return mix_trap(a, b, c, d)




def loob2bool(lb) -> bool:
	"""
	Interpret incoming variable lb as loobean. This means that true and false are swapped
	"""
	return not lb


def end(a: bloq, b: patU, c: atom):
	"""
	Tail

	Produces an atom by taking the last b blocks of size a from c.
	"""

	"""		
	++  end
	  ~/  %end
	  |=  [a=bloq b=@u c=@]
	  (mod c (bex (mul (bex a) b)))
	"""

	return c % bex(bex(a) * b)





def lsh(a: bloq, b: patU, c: atom):
	"""
	Left-shift

	Produces an atom by left-shifting c by b blocks of size a.
	"""
	"""
	 ++  lsh
		  ~/  %lsh
		  |=  [a=bloq b=@u c=@]
		  (mul (bex (mul (bex a) b)) c)
	"""
	return bex(bex(a) * b) * c



def rsh(a: bloq, b: patU, c: atom):
	"""
	Right-shift

	Right-shifts c by b blocks of size a, producing an atom.
	"""
	"""
	++  rsh
	  ~/  %rsh
	  |=  [a=bloq b=@u c=@]
	  (div c (bex (mul (bex a) b)))
	"""
	return c // bex(bex(a) * b)




def bex(a: atom) -> atom:
	"""
	Computes the result of 2^a, producing an atom.
	"""
	"""
	++  bex
	  ~/  %bex
	  |=  a=@
	  ^-  @
	  ?:  =(0 a)  1
	  (mul 2 $(a (dec a)))
	"""
	if 0 == a: return 1
	return (2 * bex(a - 1))

##############################################################################
# broken mug impl

def mug(a: noun) -> atom:
	### DOES NOT WORK LOL
	"""
	++  mug                                                 ::  mug with murmur3
		~/  %mug
		|=  a/*
		|^  (trim ?@(a a (mix $(a -.a) (mix 0x7fff.ffff $(a +.a)))))
		++  trim                                              ::  31-bit nonzero
			|=  key/@
			=+  syd=0xcafe.babe
			|-  ^-  @
			=+  haz=(muk syd (met 3 key) key)
			=+  ham=(mix (rsh 0 31 haz) (end 0 31 haz))
			?.(=(0 ham) ham $(syd +(syd)))
	  --

	"""
		
	# so after setting the max recursion depth to 100_000, 
	# python just segfaulted. either i ported this wrong or something
	# but i think for now im gonna find a third party solution		
	# update: mmh3 has hash_bytes but it will take me more than 30s to figure it out so i'm gonna do that later
	# def trim(key: atom):
		syd = 0xcafebabe
		def trim_trap(a, key, syd, haz, ham) -> atom:
			haz = muk(syd, met(3, key), key)
			print(haz)
			ham = mix(rsh(0, 31, haz), end(0, 31, haz))
			print(ham)
			if 0 == ham:
				return trim_trap(a, key, syd + 1, haz, ham)
			else:
				return ham
	
	return trim(a if type(a) == atom else \
			mix(mug(cell_head(a)), mix(0x7fffffff, mug(cell_tail(a))))
		)
	


def muk(syd: atom, len: atom, key: atom):
	"""
	++  muk
		~%  %muk  ..muk  ~
		=+  ~(. fe 5)
		|=  [syd=@ len=@ key=@]
		?>  &((lte (met 5 syd) 1) (lte (met 0 len) 31))
		=/  pad      (sub len (met 3 key))
		=/  data     (weld (rip 3 key) (reap pad 0))
		=/  nblocks  (div len 4)                             ::  intentionally off-by-one
		=/  h1  syd
		=+  [c1=0xcc9e.2d51 c2=0x1b87.3593]
		=/  blocks  (rip 5 key)
		=/  i  nblocks
		=.  h1  =/  hi  h1  |-
			?:  =(0 i)  hi                                     ::  negative array index
				=/  k1  (snag (sub nblocks i) blocks)
				=.  k1  (sit (mul k1 c1))
				=.  k1  (rol 0 15 k1)
				=.  k1  (sit (mul k1 c2))
				=.  hi  (mix hi k1)
				=.  hi  (rol 0 13 hi)
				=.  hi  (sum (sit (mul hi 5)) 0xe654.6b64)
			$(i (dec i))
		=/  tail  (slag (mul 4 nblocks) data)
		=/  k1    0
		=/  tlen  (dis len 3)
		=.  h1
			?+  tlen  h1                                       ::  fallthrough switch
				$3  =.  k1  (mix k1 (lsh 0 16 (snag 2 tail)))
					=.  k1  (mix k1 (lsh 0 8 (snag 1 tail)))
					=.  k1  (mix k1 (snag 0 tail))
					=.  k1  (sit (mul k1 c1))
					=.  k1  (rol 0 15 k1)
					=.  k1  (sit (mul k1 c2))
				(mix h1 k1)
			$2  =.  k1  (mix k1 (lsh 0 8 (snag 1 tail)))
					=.  k1  (mix k1 (snag 0 tail))
					=.  k1  (sit (mul k1 c1))
					=.  k1  (rol 0 15 k1)
					=.  k1  (sit (mul k1 c2))
				(mix h1 k1)
			$1  =.  k1  (mix k1 (snag 0 tail))
					=.  k1  (sit (mul k1 c1))
					=.  k1  (rol 0 15 k1)
					=.  k1  (sit (mul k1 c2))
				(mix h1 k1)
			==
		=.  h1  (mix h1 len)
		|^  (fmix32 h1)
		++  fmix32
			|=  h=@
				=.  h  (mix h (rsh 0 16 h))
				=.  h  (sit (mul h 0x85eb.ca6b))
				=.  h  (mix h (rsh 0 13 h))
				=.  h  (sit (mul h 0xc2b2.ae35))
				=.  h  (mix h (rsh 0 16 h))
			h
		--
	"""
	# =+  ~(. fe 5)  # todo how tf to translate?
	fe5 = fe(5)
	assert met(5, syd) <= 1 and met(0, len) <= 31
	pad = len - met(3, key)
	data = weld(rip(3, key), reap(pad, 0))
	nblocks = len // 4                             #  intentionally off-by-one
	h1 = syd
	c1 = 0xcc9e2d51
	c2 = 0x1b873593
	blocks = rip(5, key)
	i = nblocks

	def muk_trap(fe5, nblocks, k1, h1, c1, c2, blocks, i):
		hi = h1

		if (0 == i):  #  negative array index
			return hi
		else:
			k1 = snag(nblocks - i, blocks)
			k1 = fe5.sit(k1 * c1)
			k1 = fe5.rol(0, 15, k1)
			k1 = fe5.sit(k1 * c2)
			hi = mix(hi, k1)
			hi = fe5.rol(0, 13, hi)
			hi = fe5.sum(fe5.sit(hi * 5), 0xe6546b64)
			muk_trap(fe5, nblocks, k1, h1, c1, c2, blocks, i=(i - 1))

	h1 = muk_trap(fe5, nblocks, k1, h1, c1, c2, blocks, i)

	tail = slag((4 * nblocks), data)
	k1 = 0
	tlen = len & 3
	
	def calc_h1():
		nonlocal h1, c1, c2, tail, k1, fe5
		if tlen == 3:
			k1 = mix(k1, lsh(0, 16, snag(2, tail)))
			k1 = mix(k1, lsh(0, 8, snag(1, tail)))
			k1 = mix(k1, snag(0, tail))
			k1 = fe5.sit(k1 * c1)
			k1 = fe5.rol(0, 15, k1)
			k1 = fe5.sit(k1 * c2)
			return mix(h1, k1)
		elif tlen == 2:
			k1 = mix(k1, lsh(0, 8, snag(1, tail)))
			k1 = mix(k1, snag(0, tail))
			k1 = fe5.sit(k1 * c1)
			k1 = fe5.rol(0, 15, k1)
			k1 = fe5.sit(k1 * c2)
			return mix(h1, k1)
		elif tlen == 1:
			k1 = mix(k1, snag(0, tail))
			k1 = fe5.sit(k1 * c1)
			k1 = fe5.rol(0, 15, k1)
			k1 = fe5.sit(k1 * c2)
			return mix(h1, k1)
		else:
			return h1

	h1 = calc_h1()
	h1 = mix(h1, len)
	
	def fmix32(h: atom):
		h = mix(h, rsh(0, 16, h))
		h = fe5.sit(h * 0x85ebca6b)
		h = mix(h, rsh(0, 13, h))
		h = fe5.sit(h * 0xc2b2ae35)
		h = mix(h, rsh(0, 16, h))
		return h

	return fmix32(h1)

########## utility functions used in muk
def weld(a: List, b: List):
	# wet gate but everything is wet in python ;)
	# don't need calls to homo (from zuse)
	# again, just like with rip we could impl this with native python list comprehensions later
	def weld_trap() -> type(b):
		if a is None:
			return b
	return [a[0], weld(a[1:], b)]

def rip(a: bloq, b: atom) -> List[atom]:
	if 0 == b: return []
	return [end(a, 1, b), rip(b, rsh(a, 1, b))]


def snag(a: atom, b: List):
	# wet gate
	def snag_trap():
		if b is None:
			print("snag-fail")
			raise Exception()

		if 0 == a:
			return b[0]

		return snag(a - 1, b[1:])



def reap(a: atom, b: noun) -> cell:
	def reap_trap() -> List[type(b)]:
		# reap doesn't follow the style recommendation of ?~ lol. a is an atom, not true null
		#if a is None:
		#	return None

		if a == 0:
			return None
		return [b, reap(a - 1, b)]



class fe:
	"""
	modular arithmetic ring
	"""
	def __init__(self, a: bloq):
		self.a = a

	def dif(self, b: atom, c: atom) -> patS:
		"""
		++  dif  |=([b=@ c=@] (sit (sub (add out (sit b)) (sit c))))
		"""

		return self.sit((self.out + self.sit(b)) - self.sit(c))

	def inv(self, b: atom) -> atom:
		"""
		++  inv  |=(b=@ (sub (dec out) (sit b)))
		"""
		# above code uses bunt value for call to out arm
		return (self.out - 1) - self.sit(b)

	def net(self, b: atom) -> atom:
		a = self.a
		b = self.sit(b)
		if a <= 3:
			return b
		c = a - 1
		con(
				lsh(c, 1, fe(c).net(cut(c, [0, 1], b))),
				fe(c).net(cut(c, [1, 1], b))
			)

	@property
	def out(self):
		return bex(bex(self.a))

	def rol(self, b: bloq, c: atom, d: atom) -> atom:
		e = self.sit(d)
		f = bex(self.a - b)
		g = c % f
		return self.sit(con(lsh(b, g, e), (rsh(b, (f - g), e))))

	def rol(self, b: bloq, c: atom, d: atom) -> atom:
		e = self.sit(d)
		f = bex(self.a - b)
		g = c % f
		return self.sit(con(rsh(b, g, e), (lsh(b, (f - g), e))))

	def sit(self, b: atom):
		return end(self.a, 1, b)
	
	def sum(self, b: atom, c: atom):
		return self.sit(b + c)

def con(a: atom, b: atom) -> atom:
	return a | b