{-# OPTIONS_GHC -XTypeFamilies #-}

import Data.LazyNondet

type Arg a m = Typed a (Constr (Choices m))
type Res a m = UniqSupply -> Choices m -> Arg a m

false, true :: MonadPlus m => Arg Bool m
false = Typed (return (Cons 0 []))
true  = Typed (return (Cons 1 []))

guard :: MonadPlus m => Arg Bool m -> Arg a m -> Res a m
guard b x _ cs =
  caseOf cs b $ \_ b' ->
  case b' of
    Cons 1 _ -> x
    _ -> failure

nil :: MonadPlus m => Arg [a] m
nil = Typed (return (Cons 0 []))

cons :: MonadPlus m => Arg a m -> Arg [a] m -> Arg [a] m
cons x xs = Typed (return (Cons 1 [unTyped x, unTyped xs]))

insert :: MonadPlus m => Arg a m -> Arg [a] m -> Res [a] m
insert x xs is cs =
  oneOf is1 [ cons x xs
            , caseOf cs xs $ \cs' xs' ->
              case xs' of
                Cons 1 [y,ys] -> cons (Typed y) (insert x (Typed ys) is2 cs')
                _ -> failure ]
 where
  (is1,is2) = splitSupply is

permute :: MonadPlus m => Arg [a] m -> Res [a] m
permute l is cs =
  caseOf cs l $ \cs' l' ->
  case l' of
    Cons 0 _ -> nil
    Cons 1 [x,xs] -> insert (Typed x) (permute (Typed xs) is1 cs') is2 cs'
 where
  (is1,is2) = splitSupply is