A simple, readable Prolog interpreter in Haskell

Prolog.hs

{-# LANGUAGE ImportQualifiedPost, BlockArguments #-}
module Main where

import Control.Applicative
import Control.Monad
import Control.Monad.RWS
import Data.Char
import Data.Foldable
import Data.List
import Data.Set qualified as Set
import Data.Map.Strict qualified as Map
import Data.IntMap.Strict qualified as IMap
import Text.ParserCombinators.ReadP (ReadP, readP_to_S, char, string, between, satisfy, skipSpaces, sepBy, eof)
import System.IO

type Ref = Int

data Value = Free
           | Ref !Ref
           | Functor String [Ref]
           deriving (Show, Eq, Ord)

data Term = TVar String
          | TFunctor String [Term]
          deriving (Show, Eq, Ord)

data Clause = Clause { getArgs :: [Term], getBody :: [Term] }
            deriving (Show, Eq, Ord)

type Heap = IMap.IntMap Value {- map Ref to Value -}
type Program = Map.Map String [Clause]
type Locals = Map.Map String Ref

type M = RWST Program () Heap []

-- True iff x occurs as a variable in a
occurs :: Ref -> Ref -> Heap -> Bool
occurs x a heap = x == a || case heap IMap.! a of
  Ref a'         -> occurs x a' heap
  Functor _ args -> any (\arg -> occurs x arg heap) args
  _              -> False

unify :: Ref -> Ref -> M ()
unify a b = do
  heap <- get
  case (heap IMap.! a, heap IMap.! b) of
    _ | a == b -> pure () -- already unified
    (Ref a', _) -> unify a' b
    (_, Ref b') -> unify a b'
    (Free, _) | not (occurs a b heap) -> modify $ IMap.insert a (Ref b)
    (_, Free) | not (occurs b a heap) -> modify $ IMap.insert b (Ref a)
    (Functor fa argsA, Functor fb argsB) -> do
      guard $ fa == fb
      guard $ length argsA == length argsB
      zipWithM_ unify argsA argsB
    _ -> empty

-- next available index in the heap
nextIndex :: M Ref
nextIndex = gets \m -> if null m then 0 else 1 + fst (IMap.findMax m)

allocateLocals :: Clause -> M Locals
allocateLocals (Clause args body) = do
  let localsIn (TVar v) = Set.singleton v
      localsIn (TFunctor _ xs) = Set.unions (map localsIn xs)
  let locals = Set.unions (map localsIn args ++ map localsIn body)
  i <- nextIndex
  modify $ IMap.union (IMap.fromList $ zip [i..] $ replicate (Set.size locals) Free)
  pure $ Map.fromList $ zip (toList locals) [i..]

allocateTerm :: Locals -> Term -> M Ref
allocateTerm locals (TVar v) = pure $ locals Map.! v
allocateTerm locals (TFunctor n xs) = do
  args <- traverse (allocateTerm locals) xs
  i <- nextIndex
  modify $ IMap.insert i (Functor n args)
  pure i

doBody :: Locals -> [Term] -> M ()
doBody _ [] = pure ()
doBody locals (TFunctor n args:rest) = do
  vs <- traverse (allocateTerm locals) args
  doFunctor n vs
  doBody locals rest
doBody locals (TVar v:rest) = do
  val <- gets (IMap.! (locals Map.! v))
  case val of
    Functor n args -> doFunctor n args >> doBody locals rest
    _              -> error "Clause not fully instantiated"

doClause :: [Ref] -> Clause -> M ()
doClause args (Clause params body) = do
  guard $ length args == length params
  locals <- allocateLocals (Clause params body)
  let unifyArg arg param = do
        val <- allocateTerm locals param
        unify arg val
  zipWithM_ unifyArg args params
  doBody locals body

doFunctor :: String -> [Ref] -> M ()
doFunctor name args = do
  options <- reader (Map.! name)
  asum $ map (doClause args) options

-- Similar to doClause, but there are no arguments to bind
doQuery :: [Term] -> M Locals
doQuery ts = do
  locals <- allocateLocals (Clause [] ts)
  doBody locals ts
  pure locals

-- Pretty-print a term
pretty :: Heap -> Ref -> String
pretty heap x = case heap IMap.! x of
  Free         -> "_" ++ show x
  Ref x'       -> pretty heap x'
  Functor n [] -> n
  Functor n xs -> n ++ "(" ++ intercalate ", " (map (pretty heap) xs) ++ ")"

parseProgram :: String -> Maybe Program
parseQuery :: String -> Maybe [Term]
(parseProgram, parseQuery) = (parse program, parse query)
  where
    parse :: ReadP a -> String -> Maybe a
    parse p s = case readP_to_S (p <* eof) s of
      [(x, "")] -> Just x
      _         -> Nothing

    token :: ReadP a -> ReadP a
    token x = x <* skipSpaces

    openParen  = token $ string "("
    closeParen = token $ string ")"
    comma      = token $ string ","
    period     = token $ string "."
    turnstyle  = token $ string ":-"

    rigidName = token $ (:) <$> satisfy isLower <*> many (satisfy isAlphaNum)
    varName   = token $ (:) <$> (char '_' <|> satisfy isUpper) <*> many (satisfy isAlphaNum)

    var      = TVar <$> varName
    atom     = TFunctor <$> rigidName <*> pure []
    functor  = TFunctor <$> rigidName <*> between openParen closeParen (sepBy term comma)

    term = var <|> atom <|> functor

    clause = do
      TFunctor name args <- term
      body <- (turnstyle >> sepBy term comma) <|> pure []
      _ <- period
      pure (name, Clause args body)

    program = Map.fromListWith (++) . map (fmap pure) <$> many clause
    query = sepBy term comma <* period

-- A lazy list of all pretty-printed solutions for this query
runQuery :: [Term] -> Program -> [String]
runQuery q p = map prettyLocals $ runRWST (doQuery q) p IMap.empty
  where prettyLocals (locals, heap, _) =
          unlines $ map (\(name,value) -> name ++ " := " ++ pretty heap value) $ Map.toList locals

interactQuery :: [String] -> IO ()
interactQuery [] = putStrLn "no."
interactQuery (x:xs) = do
  putStrLn "yes:"
  putStr x
  l <- getLine
  when (l == ";") $ interactQuery xs

repl :: Program -> IO ()
repl p = do
  putStr "?- "
  hFlush stdout
  input <- getLine
  case stripPrefix ":load " input of
    Just f -> do
      contents <- readFile f
      case parseProgram contents of
        Just p' -> repl p'
        Nothing -> putStrLn ("Parse error in " ++ f) >> repl p
    Nothing -> case parseQuery input of
      Just q  -> interactQuery (runQuery q p) >> repl p
      Nothing -> putStrLn "Parse error" >> repl p

main :: IO ()
main = repl Map.empty