module Language.Futhark.TypeChecker.UnionFind
  ( TyVarNode,
    TyVarSol (..),
    makeTyVarNode,
    makeTyParamNode,
    find,
    getSol,
    getKey,
    assignNewSol,
    union,
    unionNewSol,
  )
where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.STRef
  ( STRef,
    modifySTRef',
    newSTRef,
    readSTRef,
    writeSTRef,
  )
import Language.Futhark (Liftedness, Loc)
import Language.Futhark.TypeChecker.Constraints
  ( CtType,
    Level,
    TyVar,
    TyVarInfo,
  )

type Type = CtType ()

-- | A (partial) solution for a type variable.
data TyVarSol
  = -- | Has been assigned this type.
    Solved Type
  | -- | Is an explicit (rigid) type parameter in the source program.
    Param Level Liftedness Loc
  | -- | Is unsolved but has this constraint.
    Unsolved (TyVarInfo ())
  deriving (Int -> TyVarSol -> ShowS
[TyVarSol] -> ShowS
TyVarSol -> String
(Int -> TyVarSol -> ShowS)
-> (TyVarSol -> String) -> ([TyVarSol] -> ShowS) -> Show TyVarSol
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TyVarSol -> ShowS
showsPrec :: Int -> TyVarSol -> ShowS
$cshow :: TyVarSol -> String
show :: TyVarSol -> String
$cshowList :: [TyVarSol] -> ShowS
showList :: [TyVarSol] -> ShowS
Show, TyVarSol -> TyVarSol -> Bool
(TyVarSol -> TyVarSol -> Bool)
-> (TyVarSol -> TyVarSol -> Bool) -> Eq TyVarSol
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TyVarSol -> TyVarSol -> Bool
== :: TyVarSol -> TyVarSol -> Bool
$c/= :: TyVarSol -> TyVarSol -> Bool
/= :: TyVarSol -> TyVarSol -> Bool
Eq)

-- | A node in the union-find graph containing information about a type
-- variable.
newtype TyVarNode s = Node (STRef s (NodeInfo s)) deriving (TyVarNode s -> TyVarNode s -> Bool
(TyVarNode s -> TyVarNode s -> Bool)
-> (TyVarNode s -> TyVarNode s -> Bool) -> Eq (TyVarNode s)
forall s. TyVarNode s -> TyVarNode s -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall s. TyVarNode s -> TyVarNode s -> Bool
== :: TyVarNode s -> TyVarNode s -> Bool
$c/= :: forall s. TyVarNode s -> TyVarNode s -> Bool
/= :: TyVarNode s -> TyVarNode s -> Bool
Eq)

data NodeInfo s
  = Link !(TyVarNode s)
  | Repr !ReprInfo

data ReprInfo = ReprInfo
  { ReprInfo -> TyVarSol
solution :: !TyVarSol,
    ReprInfo -> TyVar
key :: !TyVar
  }

-- | Create a fresh node of a type variable and return it. A fresh node
-- is in the equivalence class that contains only itself.
makeTyVarNode :: TyVar -> TyVarInfo () -> ST s (TyVarNode s)
makeTyVarNode :: forall s. TyVar -> TyVarInfo () -> ST s (TyVarNode s)
makeTyVarNode TyVar
tv TyVarInfo ()
constraint = do
  let r :: ReprInfo
r =
        ReprInfo
          { solution :: TyVarSol
solution = TyVarInfo () -> TyVarSol
Unsolved TyVarInfo ()
constraint,
            key :: TyVar
key = TyVar
tv
          }
  ref <- NodeInfo s -> ST s (STRef s (NodeInfo s))
forall a s. a -> ST s (STRef s a)
newSTRef (NodeInfo s -> ST s (STRef s (NodeInfo s)))
-> NodeInfo s -> ST s (STRef s (NodeInfo s))
forall a b. (a -> b) -> a -> b
$ ReprInfo -> NodeInfo s
forall s. ReprInfo -> NodeInfo s
Repr ReprInfo
r
  pure $ Node ref

-- | Create a fresh node of a type parameter and return it. A fresh node
-- is in the equivalence class that contains only itself.
makeTyParamNode :: TyVar -> Level -> Liftedness -> Loc -> ST s (TyVarNode s)
makeTyParamNode :: forall s. TyVar -> Int -> Liftedness -> Loc -> ST s (TyVarNode s)
makeTyParamNode TyVar
tv Int
lvl Liftedness
lft Loc
loc = do
  let r :: ReprInfo
r =
        ReprInfo
          { solution :: TyVarSol
solution = Int -> Liftedness -> Loc -> TyVarSol
Param Int
lvl Liftedness
lft Loc
loc,
            key :: TyVar
key = TyVar
tv
          }
  ref <- NodeInfo s -> ST s (STRef s (NodeInfo s))
forall a s. a -> ST s (STRef s a)
newSTRef (NodeInfo s -> ST s (STRef s (NodeInfo s)))
-> NodeInfo s -> ST s (STRef s (NodeInfo s))
forall a b. (a -> b) -> a -> b
$ ReprInfo -> NodeInfo s
forall s. ReprInfo -> NodeInfo s
Repr ReprInfo
r
  pure $ Node ref

-- | @find node@ returns the representative of @node@'s
-- equivalence class and the information associated with
-- this equivalence class.
--
-- This method performs the path compression.
find :: TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find :: forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find node :: TyVarNode s
node@(Node STRef s (NodeInfo s)
ref) = do
  node_info <- STRef s (NodeInfo s) -> ST s (NodeInfo s)
forall s a. STRef s a -> ST s a
readSTRef STRef s (NodeInfo s)
ref
  case node_info of
    -- Input node is representative.
    Repr ReprInfo
repr_info -> (TyVarNode s, ReprInfo) -> ST s (TyVarNode s, ReprInfo)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TyVarNode s
node, ReprInfo
repr_info)
    -- Input node's parent is another node.
    Link TyVarNode s
parent -> do
      a@(repr, _) <- TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
parent
      when (repr /= parent) $
        -- Performing path compression.
        writeSTRef ref $
          Link repr
      pure a

-- | Return the solution associated with the argument node's
-- equivalence class.
getSol :: TyVarNode s -> ST s TyVarSol
getSol :: forall s. TyVarNode s -> ST s TyVarSol
getSol TyVarNode s
node = ReprInfo -> TyVarSol
solution (ReprInfo -> TyVarSol)
-> ((TyVarNode s, ReprInfo) -> ReprInfo)
-> (TyVarNode s, ReprInfo)
-> TyVarSol
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TyVarNode s, ReprInfo) -> ReprInfo
forall a b. (a, b) -> b
snd ((TyVarNode s, ReprInfo) -> TyVarSol)
-> ST s (TyVarNode s, ReprInfo) -> ST s TyVarSol
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
node

-- | Return the name of the representative type variable.
getKey :: TyVarNode s -> ST s TyVar
getKey :: forall s. TyVarNode s -> ST s TyVar
getKey TyVarNode s
node = ReprInfo -> TyVar
key (ReprInfo -> TyVar)
-> ((TyVarNode s, ReprInfo) -> ReprInfo)
-> (TyVarNode s, ReprInfo)
-> TyVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TyVarNode s, ReprInfo) -> ReprInfo
forall a b. (a, b) -> b
snd ((TyVarNode s, ReprInfo) -> TyVar)
-> ST s (TyVarNode s, ReprInfo) -> ST s TyVar
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
node

-- | Assign a new solution/type to the node's equivalence class.
--
-- Precondition: The node is in an equivalence class representing an
-- unsolved/flexible type variable.
assignNewSol :: TyVarNode s -> TyVarSol -> ST s ()
assignNewSol :: forall s. TyVarNode s -> TyVarSol -> ST s ()
assignNewSol TyVarNode s
node TyVarSol
new_sol = do
  (Node ref, repr_info) <- TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
node
  modifySTRef' ref $ const . Repr $ repr_info {solution = new_sol}

-- | Join the equivalence classes of the nodes. The resulting equivalence
-- class has the same solution and key as the second argument.
union :: TyVarNode s -> TyVarNode s -> ST s ()
union :: forall s. TyVarNode s -> TyVarNode s -> ST s ()
union TyVarNode s
n1 TyVarNode s
n2 = do
  Node ref <- (TyVarNode s, ReprInfo) -> TyVarNode s
forall a b. (a, b) -> a
fst ((TyVarNode s, ReprInfo) -> TyVarNode s)
-> ST s (TyVarNode s, ReprInfo) -> ST s (TyVarNode s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
n1
  root2 <- fst <$> find n2

  writeSTRef ref $ Link root2

-- | Join the equivalence classes of the nodes. The resulting equivalence
-- class has the same key as the second argument while @new_sol@ is the
-- new solution.
unionNewSol :: TyVarNode s -> TyVarNode s -> TyVarSol -> ST s ()
unionNewSol :: forall s. TyVarNode s -> TyVarNode s -> TyVarSol -> ST s ()
unionNewSol TyVarNode s
n1 TyVarNode s
n2 TyVarSol
new_sol = do
  Node ref1 <- (TyVarNode s, ReprInfo) -> TyVarNode s
forall a b. (a, b) -> a
fst ((TyVarNode s, ReprInfo) -> TyVarNode s)
-> ST s (TyVarNode s, ReprInfo) -> ST s (TyVarNode s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TyVarNode s -> ST s (TyVarNode s, ReprInfo)
forall s. TyVarNode s -> ST s (TyVarNode s, ReprInfo)
find TyVarNode s
n1
  (root2@(Node ref2), repr_info) <- find n2

  modifySTRef' ref2 $ const . Repr $ repr_info {solution = new_sol}
  writeSTRef ref1 $ Link root2