Reasonably fast unification for Datalog atoms

DatalogUnification.hs

#!/usr/bin/env stack
-- stack script --resolver lts-14.7 --package hashable --package hashtables --package vector-sized --package transformers --ghc-options -Wall
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DataKinds #-}

module DatalogUnification (main) where

import Prelude hiding (fail)

import GHC.Generics (Generic)
import GHC.TypeLits (KnownNat)

import           Data.Hashable
import qualified Data.HashTable.Class as HT
import qualified Data.HashTable.ST.Basic as HTB
import           Data.Maybe (fromJust)
import qualified Data.Vector.Sized as V

import Control.Monad (unless, forM_)
import Control.Monad.ST
import Control.Monad.Fail
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe

newtype Literal  = Literal  String deriving (Show, Eq, Generic)
newtype Variable = Variable String deriving (Show, Eq, Generic)

data Term = Var Variable | Lit Literal deriving (Generic)

instance Hashable Literal
instance Hashable Variable
instance Hashable Term

type Substitution     = [ (Variable, Literal) ]
type SubstitutionHT s = HTB.HashTable s Variable Literal

unify :: KnownNat n
      => V.Vector n Term -> V.Vector n Literal -> Maybe Substitution
unify terms literals = runST (runMaybeT unifyM_)
  where
  unifyM_ :: MaybeT (ST s) Substitution
  unifyM_ = do
    hashTable <- lift $ HT.newSized (length terms)
    V.zipWithM_ (unifyTermM_ hashTable) terms literals
    lift $ HT.toList hashTable

  unifyTermM_ :: SubstitutionHT s -> Term -> Literal -> MaybeT (ST s) ()
  unifyTermM_ _ (Lit lit) lit'
    | lit == lit' = pure ()
    | otherwise   = fail ""
  unifyTermM_ ht (Var var) lit' = do
    mLit <- lift $ ht `HT.lookup` var
    case mLit of
      Just lit
        | lit == lit' -> pure ()
        | otherwise   -> fail ""
      Nothing -> lift $ HT.insert ht var lit'

consts1 :: V.Vector 3 Literal
consts1 = fromJust $ V.fromList [ Literal "a", Literal "b", Literal "c" ]

consts2 :: V.Vector 3 Literal
consts2 = fromJust $ V.fromList [ Literal "a", Literal "a", Literal "c" ]

atom1 :: V.Vector 3 Term
atom1 = fromJust $ V.fromList
  [ Var (Variable "X"), Var (Variable "Y"), Var (Variable "Z") ]

atom2 :: V.Vector 3 Term
atom2 = fromJust $ V.fromList
  [ Var (Variable "X"), Var (Variable "X"), Var (Variable "Z") ]

atom3 :: V.Vector 3 Term
atom3 = fromJust $ V.fromList
  [ Lit (Literal "a"), Var (Variable "Y"), Var (Variable "Z") ]

testCases :: [ (V.Vector 3 Term, V.Vector 3 Literal, Maybe Substitution) ]
testCases =
  [ ( atom1
    , consts1
    , Just
      [ (Variable "X", Literal "a")
      , (Variable "Y", Literal "b")
      , (Variable "Z", Literal "c")
      ]
    )
  , ( atom2
    , consts1
    , Nothing
    )
  , ( atom3
    , consts1
    , Just
      [ (Variable "Y", Literal "b")
      , (Variable "Z", Literal "c")
      ]
    )
  , ( atom1
    , consts2
    , Just
      [ (Variable "X", Literal "a")
      , (Variable "Y", Literal "a")
      , (Variable "Z", Literal "c")
      ]
    )
  , ( atom2
    , consts2
    , Just
      [ (Variable "X", Literal "a")
      , (Variable "Z", Literal "c")
      ]
    )
  ]

main :: IO ()
main =
  forM_ testCases $ \(atom, consts, expectation) -> do
    let result = unify atom consts
    unless (result == expectation) $ do
      putStrLn "Was expecting:"
      print result
      putStrLn "But got:"
      print expectation
      putStrLn ""
      fail "Test case failed."