{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Play where

import qualified Data.Text as T
import Control.Monad.Free
import qualified Data.Map as M
import Control.Lens hiding (cons)
import Control.Monad.State
import Control.Monad.Except
import Data.List (intercalate)
import Data.String
import Control.Applicative
import Data.Functor.Foldable (Fix(..))

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

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

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

type Expression = Fix 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) ++ "]"

newtype Var
  = Var { getVar :: T.Text }
  deriving (Eq, Ord)

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

data UnionFindPointer
  = Parent Var
  | Linked UnionFindExpr

$(makePrisms ''UnionFindPointer)

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

$(makeLenses ''UnionFindState)

type MonadUFS m = MonadState UnionFindState m

fresh :: MonadUFS m => m Var
fresh = do
  ufsSupply %= tail
  use (ufsSupply . to head)

find :: MonadUFS m => Var -> m (Var, Maybe UnionFindExpr)
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!"

data UnificationError
  = AtomMismatch Identifier Identifier
  | AtomNotCons Identifier Var [Var]
  | ConsLengthMismatch UnionFindExpr UnionFindExpr
  deriving (Show)

unify :: (MonadError UnificationError m, MonadUFS 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) -> go p q where
      go (Atom x) (Atom y)
        | getIdentifier x == getIdentifier y = pure ()
        | otherwise = throwError (AtomMismatch x y)
      go (Atom x) (Cons y ys) = throwError (AtomNotCons x y ys)
      go m@Cons{} n@Atom{} = go n m
      go m@(Cons x xs) n@(Cons y ys)
        | length xs == length ys = do
          unify x y
          zipWithM_ unify xs ys
          ufsMap . at xu .= Just (Parent xv)
        | otherwise = throwError (ConsLengthMismatch m n)

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

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

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

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 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 :: (MonadError UnificationError m, MonadUFS 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.