pyrites/haskell-experiments/src/Datalog/NaiveDatabase.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications  #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Redundant flip" #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Datalog.NaiveDatabase where

import Data.Map (Map)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Datalog.DatalogParser(parseDatalog, Literal(..), Statement(..), Term (..), Head (HeadSingle))
import qualified Data.Map as Map
import Text.Megaparsec (ParseErrorBundle)
import Data.Void
import Control.Exception.Base
import Data.List
import Data.Maybe

data Value =
  ValueInt Int |
  ValueSymbol String

data NaiveDatabase = NaiveDatabase {
  relations :: Map RelationId Relation,
  constants :: Set Constant
}

data RuleElement = -- entry occurring in a head or body relation - constant or variable
  RuleElementConstant Constant |
  RuleElementVariable Text
  deriving (Show, Eq)

data RelationRule = RelationRule {
  headVariables :: [Text],
  body :: [(Relation, [RuleElement])]
} deriving (Show, Eq)

data Relation = Relation {
  _name :: Text,
  _arity :: Int,
  _tuples :: Set [Constant],
  _rules :: [RelationRule]
} deriving (Show, Eq)

-- newtype RelationId = RelationId Text
--   deriving (Eq, Ord, Show)

-- Our constants will be the terms of the Datalog grammar - ints/variables/symbols
type Constant = Term
type RelationId = Text

-- newtype Constant = Constant Text
--   deriving (Eq, Ord, Show)

emptyDB :: NaiveDatabase
emptyDB = NaiveDatabase {
  relations = Map.empty,
  constants = Set.empty   -- the Herbrand universe
}

withFacts :: [Text] -> NaiveDatabase
withFacts facts =
  foldl addFact emptyDB (extractFact <$> facts) where
    extractFact:: Text -> Literal
    extractFact factText =
      case (parseDatalog factText) of
        Right (Fact fact) -> fact
        Right otherStatement -> throw $ NonFactException factText otherStatement
        Left ex -> throw $ CannotParseStatementException factText ex
    addFact :: NaiveDatabase -> Literal -> NaiveDatabase
    addFact (NaiveDatabase relationMap constantSet) (Literal neg relationName terms) =
      NaiveDatabase newRelationMap newConstantSet where
        newArity = length terms
        newRelation = lookupRelation relationName relationMap newArity terms
        newRelationMap = Map.insert relationName newRelation relationMap
        newConstantSet = Set.union constantSet $ Set.fromList terms

lookupRelation :: RelationId -> Map RelationId Relation -> Int -> [Term] -> Relation
lookupRelation relationName relationMap newArity terms =
    case Map.lookup relationName relationMap of
      Nothing -> Relation relationName newArity (Set.singleton terms) []
      Just relation ->
        if (_arity relation == newArity)
          then Relation relationName (length terms) (Set.singleton terms) []
          else throw $ BadArityException relationName newArity

data ConstraintElement = -- entry occurring in a rule body constraint - constant or variable index
  ConstraintElementConstant Constant |
  ConstraintElementIndex Int
  deriving (Show, Eq)

data BodyConstraint = BodyConstraint {
  _relation :: Relation,
  _elements :: [ConstraintElement]
}

data RuleContext = RuleContext {
  __relation :: Relation,
  _headVariables :: [RuleElement],
  _bodyConstraints :: [BodyConstraint],
  _db :: NaiveDatabase
}

toElement :: Term -> RuleElement
-- toElement constant@(Var _) = RuleElementConstant constant
-- toElement constant@(Num _) = RuleElementConstant constant
toElement (Var name) = RuleElementVariable name
toElement constant = RuleElementConstant constant

maybeConstant :: RuleElement -> Maybe Constant
maybeConstant (RuleElementConstant constant) = Just constant
maybeConstant _ = Nothing

withFactsAndRules :: [Text] -> [Text] -> NaiveDatabase
withFactsAndRules facts rules =
  foldl addRule (withFacts facts) (extractRule <$> rules) where
    extractRule:: Text -> (Literal, [Literal])
    extractRule ruleText =
      case (parseDatalog ruleText) of
        Right (Rule (HeadSingle lhs) rhs) -> (lhs, rhs)
        Right otherStatement -> throw $ NonRuleException ruleText otherStatement
        Left ex -> throw $ CannotParseStatementException ruleText ex
    digestHead :: RuleContext -> Literal -> RuleContext
    digestHead  context (Literal neg relationName terms) =
      RuleContext { 
        __relation = relation, 
        _headVariables = variables', 
        _bodyConstraints = _bodyConstraints context, 
        _db = NaiveDatabase relationMap' constants'
      } where
        newArity = length terms
        relation = __relation context
        relationMap :: Map RelationId Relation = (relations (_db context))
        relationMap' = Map.insert relationName relation relationMap
        extraVariables = toElement <$> terms
        extraConstants = catMaybes $ maybeConstant <$> extraVariables
        variables' = nub $ _headVariables context ++ extraVariables
        constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
    digestBody :: RuleContext -> Literal -> RuleContext
    digestBody context (Literal neg subRelationName terms) =
      RuleContext { 
        __relation = relation, 
        _headVariables = variables', 
        _bodyConstraints = constraints', 
        _db = NaiveDatabase relationMap' constants'
      }  where
        relation = __relation context
        newArity = length terms
        subRelation = lookupRelation subRelationName relationMap newArity terms
        relationMap :: Map RelationId Relation = relations (_db context)
        relationMap' = Map.insert subRelationName subRelation relationMap
        extraVariables = toElement <$> terms
        extraConstants = catMaybes $ maybeConstant <$> extraVariables
        variables' = nub $ _headVariables context ++ extraVariables
        constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
        constraints = _bodyConstraints context
        newConstraint = BodyConstraint subRelation (toConstraint <$> terms)
        constraints' = constraints ++ [newConstraint]
        varIndex :: Text -> Int
        varIndex name =
          case elemIndex (RuleElementVariable name) variables' of
            Just index -> index
            Nothing -> throw $ VariableLookupException name variables'
        toConstraint :: Term -> ConstraintElement
        -- toConstraint thing = ConstraintElementIndex 0
        toConstraint (Var name) = ConstraintElementIndex (varIndex name)
        toConstraint constant = ConstraintElementConstant constant
    addRule :: NaiveDatabase -> (Literal, [Literal]) -> NaiveDatabase
    addRule  db (ruleHead, body) =
      NaiveDatabase newRelationMap constants' where
        relationName = predName ruleHead
        terms = arguments ruleHead
        newArity = length terms
        relation = lookupRelation relationName (relations (_db context)) newArity terms
        context = RuleContext { 
          __relation = relation, 
          _headVariables = [], 
          _bodyConstraints = [], 
          _db = db
        }
        context' = digestHead context ruleHead
        context'' = foldl digestBody context' body
        db' = _db context''
        relationMap = relations db'
        variables' = _headVariables context''
        varIndex :: Text -> Int -- TODO unify with the above
        varIndex name =
          case elemIndex (RuleElementVariable name) variables' of
            Just index -> index
            Nothing -> throw $ VariableLookupException name variables'
        toConstraintElement :: RuleElement -> ConstraintElement
        toConstraintElement (RuleElementConstant constant) = ConstraintElementConstant constant
        toConstraintElement (RuleElementVariable varName) = ConstraintElementIndex (varIndex varName)
        extractVarName :: RuleElement -> Text
        extractVarName (RuleElementVariable varName) = varName
        extractVarName (RuleElementConstant constant) = throw $ UnexpectedConstantException constant
        -- toBodyConstraint :: (Relation, [RuleElement]) -> BodyConstraint
        -- toBodyConstraint (subRelation, ruleElements) =
        --   BodyConstraint {
        --     _relation = subRelation,
        --     _elements = toConstraintElement <$> ruleElements
        --   }
        newRule = RelationRule {
          headVariables = extractVarName <$> variables',
          body = fromBodyConstraint <$> (_bodyConstraints context'')
        } where
          fromBodyConstraint :: BodyConstraint -> (Relation, [RuleElement])
          fromBodyConstraint (BodyConstraint subRelation elements) =
            (subRelation, toRuleElement <$> elements) 
          toRuleElement :: ConstraintElement -> RuleElement
          toRuleElement (ConstraintElementConstant constant) = RuleElementConstant constant
          toRuleElement (ConstraintElementIndex index) = variables' !! index
        -- input is [(Relation, [RuleElement])]
        -- bodyConstraint has
        --   _relation :: Relation,
        --   _elements :: [ConstraintElement]
        relation' = Relation {
          _name = _name relation,
          _arity = newArity,
          _tuples = _tuples relation,
          _rules = (_rules relation) ++ [newRule]
        }
        newRelationMap = Map.insert relationName relation' relationMap
        constants' = constants db'

query :: NaiveDatabase -> Text -> Text
query db qText =
  case (parseDatalog qText) of
    Right (Query texts literals) -> "#NYI"
    Right otherStatement -> throw $ NonQueryException qText otherStatement
    Left ex -> throw $ CannotParseStatementException qText ex

data NaiveDatabaseException
  = CannotParseStatementException Text (ParseErrorBundle Text Void) |
    NonFactException Text Statement |
    NonRuleException Text Statement |
    NonQueryException Text Statement |
    BadArityException Text Int |
    VariableLookupException Text [RuleElement] |
    UnexpectedConstantException Constant
  deriving (Show)

instance Exception NaiveDatabaseException