ingesting rules correctly

haskell-experiments/src/Datalog/NaiveDatabase.hs

@@ -4,6 +4,7 @@
 
 {-# HLINT ignore "Redundant flip" #-}
 {-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 module Datalog.NaiveDatabase where
 
 import Data.Map (Map)
@@ -11,12 +12,14 @@ import Data.Set (Set)
 import Data.Set qualified as Set
 import Data.Text (Text)
 import qualified Data.Text as T
-import Datalog.DatalogParser(parseDatalog, Literal(Literal), Statement(..), Term, Head (HeadSingle))
+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 (Exception)
 import Control.Exception.Base
+import Data.List
+import Data.Maybe
 
 data Value =
   ValueInt Int |
@@ -27,20 +30,20 @@ data NaiveDatabase = NaiveDatabase {
   constants :: Set Constant
 }
 
-data BodyEntry = -- entry occurring in a head or body relation - constant or variable
+data RuleElement = -- entry occurring in a head or body relation - constant or variable
-  BodyEntryConstant Constant |
+  RuleElementConstant Constant |
-  BodyEntryVariable Text
+  RuleElementVariable Text
   deriving (Show, Eq)
 
 data RelationRule = RelationRule {
-  headVariables :: [BodyEntry],
+  headVariables :: [RuleElement],
-  body :: [(Relation, [BodyEntry])]
+  body :: [(Relation, [RuleElement])]
 } deriving (Show, Eq)
 
 data Relation = Relation {
-  arity :: Int,
+  _arity :: Int,
-  tuples :: Set [Constant],
+  _tuples :: Set [Constant],
-  rules :: [RelationRule]
+  _rules :: [RelationRule]
 } deriving (Show, Eq)
 
 -- newtype RelationId = RelationId Text
@@ -69,18 +72,48 @@ withFacts facts =
         Right otherStatement -> throw $ NonFactException factText otherStatement
         Left ex -> throw $ CannotParseStatementException factText ex
     addFact :: NaiveDatabase -> Literal -> NaiveDatabase
-    addFact (NaiveDatabase relations constants) (Literal neg relationName terms) =
+    addFact (NaiveDatabase relationMap constantSet) (Literal neg relationName terms) =
-      NaiveDatabase newRelations newConstants where
+      NaiveDatabase newRelationMap newConstantSet where
         newArity = length terms
-        newRelation =
+        newRelation = lookupRelation relationName relationMap newArity terms
-          case Map.lookup relationName relations of
+        newRelationMap = Map.insert relationName newRelation relationMap
-            Nothing -> Relation (length terms) (Set.singleton terms) []
+        newConstantSet = Set.union constantSet $ Set.fromList terms
-            Just relation ->
+
-              if (arity relation == newArity)
+lookupRelation :: RelationId -> Map RelationId Relation -> Int -> [Term] -> Relation
-                then Relation (length terms) (Set.singleton terms) []
+lookupRelation relationName relationMap newArity terms =
-                else throw $ BadArityException relationName newArity
+    case Map.lookup relationName relationMap of
-        newRelations = Map.insert relationName newRelation relations 
+      Nothing -> Relation newArity (Set.singleton terms) []
-        newConstants = Set.union constants $ Set.fromList terms
+      Just relation ->
+        if (_arity relation == newArity)
+          then Relation (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 =
@@ -91,12 +124,106 @@ withFactsAndRules facts rules =
         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 (NaiveDatabase relations constants) ((Literal neg relationName terms), body) =
+    addRule  db (ruleHead, body) =
-      NaiveDatabase newRelations newConstants where
+      NaiveDatabase newRelationMap constants' where
-        newRelations = relations 
+        relationName = predName ruleHead
-        newConstants = constants
+        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)
+
+        -- toBodyConstraint :: (Relation, [RuleElement]) -> BodyConstraint
+        -- toBodyConstraint (subRelation, ruleElements) =
+        --   BodyConstraint {
+        --     _relation = subRelation,
+        --     _elements = toConstraintElement <$> ruleElements
+        --   }
+        newRule = RelationRule {
+          headVariables = 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 {
+          _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
@@ -109,7 +236,8 @@ data NaiveDatabaseException
     NonFactException Text Statement |
     NonRuleException Text Statement |
     NonQueryException Text Statement |
-    BadArityException Text Int
+    BadArityException Text Int |
+    VariableLookupException Text [RuleElement]
   deriving (Show)
 
 instance Exception NaiveDatabaseException

haskell-experiments/test/Test/Datalog/NaiveDatabaseSpec.hs

@@ -45,9 +45,58 @@ spec = do
         (Set.fromList $ Sym <$> ["alice", "bob", "carol"])
       relations db `shouldBe`
         Map.fromList [
-          ("parent",  
+          ("ancestor",
-            Relation 2 (Set.fromList ([Sym <$> ["bob", "carol"]])) [] )
+            Relation {
-        ]
+              _arity = 2, 
+              _tuples = Set.fromList [
+                [Var "X", Var "Y"]
+              ],
+              _rules = [
+                RelationRule {
+                  headVariables = [
+                    RuleElementVariable "X",
+                    RuleElementVariable "Y",
+                    RuleElementVariable "Z"
+                  ], 
+                  body = [
+                    (
+                      Relation {
+                        _arity = 2, 
+                        _tuples = Set.fromList [
+                          [Var "X",Var "Z" ]
+                        ], 
+                        _rules = []
+                      }, [
+                          RuleElementVariable "X",
+                          RuleElementVariable "Z"
+                        ]
+                    ),(
+                      Relation {
+                        _arity = 2, 
+                        _tuples = Set.fromList [
+                          [Var "Z",Var "Y"]
+                        ], 
+                        _rules = []
+                      },[
+                        RuleElementVariable "Z",
+                        RuleElementVariable "Y"
+                        ]
+                      )
+                    ]
+                  }
+                ]
+              }
+            ),
+            ("parent",
+              Relation {
+                _arity = 2, 
+                _tuples = Set.fromList [
+                  [Var "X",Var "Z"]
+                ], 
+                _rules = []
+              }
+            )
+          ]
 
     it "can do basic queries" $ do
       let db = NaiveDatabase.withFacts