DatalogProgram.hs: Broken visitor

The expression visitor implementation, which is supposed to visit all
expressions in the program, ignored primary key expressions.

Author: ryzhyk

src/Language/DifferentialDatalog/DatalogProgram.hs

@@ -65,6 +65,8 @@ import Language.DifferentialDatalog.Type
 -- | Map function 'fun' over all expressions in a program
 progExprMapCtxM :: (Monad m) => DatalogProgram -> (ECtx -> ENode -> m Expr) -> m DatalogProgram
 progExprMapCtxM d fun = do
+    rels' <- M.fromList <$>
+             (mapM (\(rname, rel) -> (rname,) <$> relExprMapCtxM fun rel) $ M.toList $ progRelations d)
     funcs' <- mapM (\f -> do e <- case funcDef f of
                                        Nothing -> return Nothing
                                        Just e  -> Just <$> exprFoldCtxM fun (CtxFunc f) e
@@ -75,8 +77,16 @@ progExprMapCtxM d fun = do
                              return r{ruleLHS = lhs, ruleRHS = rhs})
                    $ progRules d
     return d{ progFunctions = funcs'
+            , progRelations = rels'
             , progRules     = rules'}
 
+relExprMapCtxM :: (Monad m) => (ECtx -> ENode -> m Expr) -> Relation -> m Relation
+relExprMapCtxM fun rel = do
+    pkey' <- mapM (\key -> do e' <- exprFoldCtxM fun (CtxKey rel) $ keyExpr key
+                              return key{keyExpr = e'})
+             $ relPrimaryKey rel
+    return rel{relPrimaryKey = pkey'}
+
 atomExprMapCtxM :: (Monad m) => (ECtx -> ENode -> m Expr) -> ECtx -> Atom -> m Atom
 atomExprMapCtxM fun ctx a = do
     v <- exprFoldCtxM fun ctx $ atomVal a

test/datalog_tests/constr.fail.ast.expected

@@ -20,4 +20,6 @@ expecting function name
 
 error: ./test/datalog_tests/constr.fail.dl:4:15-7:1: Field x is declared with different types
 
-error: ./test/datalog_tests/constr.fail.dl:4:1-6:1: The following type variables are not used in type definition: D
+error: ./test/datalog_tests/constr.fail.dl:4:1-7:1: The following type variables are not used in type definition: D
+
+./test/datalog_tests/constr.fail.dl:6:54-7:1: Unknown module vm.  Did you forget to import it?

test/datalog_tests/constr.fail.dl

@@ -33,3 +33,10 @@ typedef Alt = C0{x: bit<16>}
 // type argument 'D not used in type definition
 typedef Parameterized<'A,'B,'C,'D> = Option1{x: 'A}
                                    | Option2{y: 'B, z: 'C}
+
+//---
+
+typedef VIP  = string
+typedef UUID = string
+
+input relation VM_Host[(VIP, UUID)] primary key (vm) vm.VIP