# -*- coding: utf-8 -*-

"""Nondeterministic evaluation (a tuple comprehension with a multi-expr body)."""

from ast import Tuple, arg

from mcpyrate.quotes import macros, q, u, n, a, h # noqa: F401

from mcpyrate.splicing import splice_expression

from .letdoutil import isenvassign, UnexpandedEnvAssignView

from ..amb import monadify

from ..amb import insist, deny # for re-export only # noqa: F401

from ..misc import namelambda

def forall(exprs):

"""[syntax, expr] Nondeterministic evaluation.

Fully based on AST transformation, with real lexical variables.

Like Haskell's do-notation, but here specialized for the List monad.

Example::

# pythagorean triples

pt = forall[z << range(1, 21), # hypotenuse

x << range(1, z+1), # shorter leg

y << range(x, z+1), # longer leg

insist(x*x + y*y == z*z),

(x, y, z)]

assert tuple(sorted(pt)) == ((3, 4, 5), (5, 12, 13), (6, 8, 10),

(8, 15, 17), (9, 12, 15), (12, 16, 20))

"""

if type(exprs) is not Tuple: # pragma: no cover, let's not test macro expansion errors.

raise SyntaxError("forall body: expected a sequence of comma-separated expressions")

itemno = 0

def build(lines, tree):

if not lines:

return tree

line, *rest = lines

if isenvassign(line): # no need for "let"; we just borrow a very small part of its syntax machinery.

view = UnexpandedEnvAssignView(line)

k, v = view.name, view.value

else:

k, v = "_ignored", line

islast = not rest

# don't unpack on last line to allow easily returning a tuple as a result item

Mv = q[h[monadify](a[v], u[not islast])]

if not islast:

lam = q[lambda _: n["_here_"]]

lam.args.args = [arg(arg=k)]

nonlocal itemno

itemno += 1

label = "item{itemno}" if k == "_ignored" else k

namedlam = q[h[namelambda](u[f"forall_{label}"])(a[lam])]

body = q[a[Mv] >> a[namedlam]] # monadic bind: >>

else:

body = Mv

if tree:

newtree = splice_expression(body, tree, "_here_")

else:

newtree = body

return build(rest, newtree)

return q[h[tuple](a[build(exprs.elts, None)])]