Play.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}

module Play
  ( ExactZip(..), MonadUFS, Var
  , fresh, record, report, unify
  , run
  ) where

import qualified Data.Text as T
import Control.Monad.Free
import qualified Data.Map as M
import Lens.Micro
import Lens.Micro.Mtl
import Lens.Micro.GHC () -- Instances only
import Control.Monad.State
import Control.Applicative
import Data.Foldable hiding (find)
import GHC.Generics

class ExactZip f where
  exactZip :: f a -> f b -> Maybe (f (a, b))
  default exactZip :: (Generic1 f, ExactZip (Rep1 f)) => f a -> f b -> Maybe (f (a, b))
  exactZip x y = to1 <$> exactZip (from1 x) (from1 y)

instance ExactZip V1 where
  exactZip _ _ = error "exactZip on void type?"

instance ExactZip U1 where
  exactZip U1 U1 = Just U1

instance (ExactZip f, ExactZip g) => ExactZip (f :+: g) where
  exactZip (L1 x) (L1 y) = L1 <$> exactZip x y
  exactZip (R1 x) (R1 y) = R1 <$> exactZip x y
  exactZip _ _ = Nothing

instance (ExactZip f, ExactZip g) => ExactZip (f :*: g) where
  exactZip (x1 :*: y1) (x2 :*: y2)
    = liftA2 (:*:) (exactZip x1 x2) (exactZip y1 y2)

instance Eq c => ExactZip (K1 i c) where
  exactZip (K1 a) (K1 b)
    | a == b = Just (K1 a)
    | otherwise = Nothing

instance ExactZip f => ExactZip (M1 i c f) where
  exactZip (M1 u) (M1 v) = M1 <$> exactZip u v

instance ExactZip Par1 where
  exactZip (Par1 u) (Par1 v) = Just (Par1 (u, v))

instance ExactZip f => ExactZip (Rec1 f) where
  exactZip (Rec1 u) (Rec1 v) = Rec1 <$> exactZip u v

instance ExactZip []

newtype Var
  = Var T.Text
  deriving (Eq, Ord)

instance Show Var where
  show (Var v) = "#" ++ T.unpack v

data UnionFindPointer f
  = Parent Var
  | Linked (f Var)

data UnionFindState f
  = UnionFindState
    { _ufsSupply :: [Var]
    , _ufsMap :: M.Map Var (UnionFindPointer f)
    }

ufsMap :: Lens' (UnionFindState f) (M.Map Var (UnionFindPointer f))
ufsMap f_aj0N (UnionFindState x_aj0O x_aj0P)
  = fmap (\ y_aj0Q -> UnionFindState x_aj0O y_aj0Q) (f_aj0N x_aj0P)
{-# INLINE ufsMap #-}

ufsSupply :: Lens' (UnionFindState f) [Var]
ufsSupply f_aj0R (UnionFindState x_aj0S x_aj0T)
  = fmap (\ y_aj0U -> UnionFindState y_aj0U x_aj0T) (f_aj0R x_aj0S)
{-# INLINE ufsSupply #-}

type MonadUFS f m = MonadState (UnionFindState f) m

fresh :: MonadUFS f m => m Var
fresh = do
  ufsSupply %= tail
  head <$> use ufsSupply

find :: MonadUFS f m => Var -> m (Var, Maybe (f Var))
find u = locateRoot u >>= \(x, e) -> (x, e) <$ compressPath u x where
  locateRoot t =
    use (ufsMap . at t) >>= \case
      Just (Parent v) -> locateRoot v
      Just (Linked ex) -> pure (t, Just ex)
      Nothing -> pure (t, Nothing)
  compressPath t x = go t where
    go m
      | m == x = pure ()
      | otherwise = use (ufsMap . at m) >>= \case
          Just (Parent par) -> do
            ufsMap . at m .= Just (Parent x)
            go par
          _ -> error "Internal error: find: Can't happen!"

unify :: (Foldable f, ExactZip f, Alternative m, MonadUFS f m) => Var -> Var -> m ()
unify u1 v1 = do
  (xu, eu) <- find u1
  (xv, ev) <- find v1
  when (xu /= xv) $ case (eu, ev) of
    (Nothing, _) -> ufsMap . at xu .= Just (Parent xv)
    (Just _, Nothing) -> ufsMap . at xv .= Just (Parent xu)
    (Just p, Just q) -> case exactZip p q of
      Nothing -> empty
      Just ex -> do
        for_ ex (uncurry unify)
        ufsMap . at xu .= Just (Parent xv)

record :: (Traversable f, MonadUFS f m) => Free f Var -> m Var
record = iterA go where
  go f = do
    u <- sequence f
    v <- fresh
    ufsMap . at v .= Just (Linked u)
    pure v

report :: (Traversable f, MonadUFS f m) => Var -> m (Free f Var)
report = unfoldM go where
  go v = find v >>= \case
    (u, Nothing) -> pure (Left u)
    (_, Just t) -> pure (Right t)

run :: StateT (UnionFindState f) Maybe a
    -> Maybe a
run s = evalStateT s (UnionFindState vars M.empty) where
  vars =
    let u = "" : liftA2 (flip (:)) u ['A'..'Z']
    in map (Var . T.pack) u

Expr.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}

module Expr where

import Data.String
import Control.Monad.Free
import Data.List
import qualified Data.Text as T
import GHC.Generics

import Control.Applicative
import Play

newtype Identifier
  = Identifier { getIdentifier :: T.Text }
  deriving (IsString, Eq)

instance Show Identifier where
  show (Identifier u) = T.unpack u

data ExprF a
  = Atom Identifier
  | Cons a [a]
  deriving (Functor, Foldable, Traversable, Generic1)

instance ExactZip ExprF

type PartialExpr = Free ExprF Var
type UnionFindExpr = ExprF Var

printPartial :: PartialExpr -> String
printPartial = iter go . fmap show where
  go (Atom d) = show d
  go (Cons x xs) = x ++ "[" ++ intercalate "," xs ++ "]"

instance Show u => Show (ExprF u) where
  show (Atom (Identifier u)) = T.unpack u
  show (Cons x xs) = show x ++ "[" ++ intercalate "," (show <$> xs) ++ "]"

atom :: Identifier -> PartialExpr
atom = Free . Atom

cons :: PartialExpr -> [PartialExpr] -> PartialExpr
cons u v = Free (Cons u v)

instance IsString PartialExpr where
  fromString = atom . fromString

a1, a2 :: MonadUFS f m => m PartialExpr

-- f[#A, u[#B], #C]
a1 = do
  a <- fresh
  b <- fresh
  c <- fresh
  pure $ cons "f" [pure a, cons "u" [pure b], pure c]

-- #D[#E, #F, #G[v]]
a2 = do
  d <- fresh
  e <- fresh
  f <- fresh
  g <- fresh
  pure $ cons (pure d) [pure e, pure f, cons (pure g) ["v"]]

test :: (Alternative m, MonadUFS ExprF m) => m PartialExpr
test = do
  e1 <- a1
  e2 <- a2
  x <- record e1
  y <- record e2
  unify x y
  report x

-- ghci> let Right u = run (printPartial <$> test) in putStrLn u
-- f[#E,u[#B],#G[v]]
--
-- Note: The variable names are generated automatically. As you can see
-- the definitions of a1 and a2 do not specify names.