tests green

haskell-experiments/src/Datalog/NaiveDatabase.hs

@@ -105,10 +105,6 @@ data RuleContext = RuleContext
   , _db :: NaiveDatabase
   }
 
-maybeConstant :: RuleElement -> Maybe Constant
-maybeConstant (RuleElementConstant constant) = Just constant
-maybeConstant _ = Nothing
-
 withFactsAndRules :: [Text] -> [Text] -> NaiveDatabase
 withFactsAndRules facts = foldr (addRule . extractRule) (withFacts facts)
  where
@@ -139,37 +135,51 @@ withFactsAndRules facts = foldr (addRule . extractRule) (withFacts facts)
     relationMap' = Map.insert relationName relation relationMap
     extraVariables = headTermToElement <$> terms
     entries' = nub extraVariables
-    extraConstants = mapMaybe maybeConstant entries'
+    extraConstants = mapMaybe extractConstant entries' where
+      extractConstant :: RuleElement -> Maybe Constant
+      extractConstant (RuleElementConstant constant) = Just constant
+      extractConstant _ = Nothing
     constants' = Set.union (constants db) $ Set.fromList extraConstants
   digestBody :: Literal -> RuleContext -> RuleContext
-  digestBody (Literal neg subRelationName terms) context =
+  digestBody (Literal neg subRelationName subTerms) context =
     RuleContext
-      { __relation = relation
+      { __relation = __relation context
       , _variableNames = variableNames
-      , _headEntries = variables'
-      , _bodyConstraints = constraints'
+      , _headEntries = _headEntries context
+      , _bodyConstraints = newConstraint : constraints
       , _db = NaiveDatabase relationMap' constants'
       }
    where
-    relation = __relation context
-    variableNames = _variableNames context
-    bodyTermToElement :: Term -> RuleElement
-    bodyTermToElement (Var name) = 
-      case (elemIndex name) variableNames of
-          Just index -> RuleElementVariable index
-          Nothing -> throw $ VariableLookupException name variableNames
-    bodyTermToElement constant = RuleElementConstant constant
-    newArity = length terms
+    variableNames = nub $ _variableNames context ++ extractVariableNames subTerms
+    newArity = length subTerms
     subRelation = lookupRelation subRelationName relationMap newArity Set.empty
     relationMap :: Map RelationId Relation = relations (_db context)
     relationMap' = Map.insert subRelationName subRelation relationMap
-    extraVariables = bodyTermToElement <$> terms
-    extraConstants = mapMaybe maybeConstant extraVariables
-    variables' = nub $ _headEntries context ++ extraVariables
+    -- extraVariables = bodyTermToElement <$> terms where
+    --   bodyTermToElement :: Term -> RuleElement
+    --   bodyTermToElement (Var name) = 
+    --     case (elemIndex name) variableNames of
+    --         Just index -> RuleElementVariable index
+    --         Nothing -> throw $ VariableLookupException name variableNames
+    --   bodyTermToElement constant = RuleElementConstant constant
+    extraConstants = mapMaybe constantFromTerm subTerms where
+      constantFromTerm :: Term -> Maybe Constant
+      constantFromTerm (Var _) = Nothing
+      constantFromTerm constant = Just constant
+    -- variables' = nub $ _headEntries context ++ extraVariables
     constants' = Set.union (constants (_db context)) $ Set.fromList extraConstants
     constraints = _bodyConstraints context
-    newConstraint = BodyConstraint subRelation variables'
-    constraints' = newConstraint : constraints
+    newConstraint = BodyConstraint subRelation subRuleElements where
+      subRuleElements = toRuleElement <$> subTerms
+      toRuleElement :: Term -> RuleElement
+      toRuleElement (Var varName) = RuleElementVariable $ lookupVariable varName
+      toRuleElement constant = RuleElementConstant constant
+      lookupVariable :: Text -> Int
+      lookupVariable varName = 
+        case (elemIndex varName) variableNames of
+            Just index -> index
+            Nothing -> throw $ VariableLookupException varName variableNames
+
   addRule :: (Literal, [Literal]) -> NaiveDatabase -> NaiveDatabase
   addRule (ruleHead, body) db  =
     NaiveDatabase relationMap' constants'