Make ceil generation of in_keys more robust (#429)

We were previously making the assumption that the only in_keys functions
are those for the default built-in maps or cell maps. However,
@virgil-serbanuta has a project which defines custom maps with a custom
in_keys hook. In order to find the correct in_key symbol, we need to
extract the "hook" attribute when internalising symbols and build a
rverse lookup map for the given key and map sorts

---------

Co-authored-by: github-actions <github-actions@github.com>
Co-authored-by: Georgy Lukyanov <georgiylukjanov@gmail.com>
Co-authored-by: Georgy Lukyanov <mail@geo2a.info>

library/Booster/Definition/Attributes/Base.hs

@@ -201,6 +201,7 @@ data SymbolAttributes = SymbolAttributes
     , hasEvaluators :: Flag "canBeEvaluated"
     , collectionMetadata :: Maybe KCollectionMetadata
     , smt :: Maybe SMTType
+    , hook :: Maybe Text
     }
     deriving stock (Eq, Ord, Show, Generic, Data, Lift)
     deriving anyclass (NFData, Hashable)

library/Booster/Definition/Attributes/Reader.hs

@@ -188,6 +188,7 @@ instance HasAttributes ParsedSymbol where
             <*> hasConcreteEvaluators
             <*> pure Nothing
             <*> smt
+            <*> (attributes .:? "hook")
 
 instance HasAttributes ParsedSort where
     type Attributes ParsedSort = SortAttributes

library/Booster/Definition/Ceil.hs

@@ -24,7 +24,7 @@ import Control.Monad.IO.Class (MonadIO (liftIO))
 import Control.Monad.Logger (MonadLoggerIO)
 import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.Writer (runWriterT, tell)
-import Data.ByteString.Char8 (isPrefixOf, stripPrefix)
+import Data.ByteString.Char8 (isPrefixOf)
 import Data.Coerce (coerce)
 import Data.Foldable (toList)
 import Data.Map qualified as Map
@@ -156,7 +156,7 @@ computeCeil (AndTerm l r) = concatMapM computeCeil [l, r]
 computeCeil (Injection _ _ t) = computeCeil t
 computeCeil (KMap _ keyVals rest) = do
     recArgs <- concatMapM computeCeil $ concat [[k, v] | (k, v) <- keyVals] <> maybeToList rest
-    symbols <- (.definition.symbols) <$> getConfig
+    symbols <- mkInKeysMap . (.definition.symbols) <$> getConfig
     pure $
         [Left $ Predicate $ mkNeq a b | a <- map fst keyVals, b <- map fst keyVals, a /= b]
             <> [ Left $ Predicate $ NotBool (mkInKeys symbols a rest') | a <- map fst keyVals, rest' <- maybeToList rest
@@ -183,33 +183,26 @@ mkNeq a b
     | sortOfTerm a == SortInt = NEqualsInt a b
     | otherwise = NEqualsK (KSeq (sortOfTerm a) a) (KSeq (sortOfTerm b) b)
 
-mkInKeys :: Map.Map SymbolName Symbol -> Term -> Term -> Term
-mkInKeys symbols k m
-    | sortOfTerm k == SortKItem && sortOfTerm m == SortMap =
-        SymbolApplication
-            ( Symbol
-                "Lbl'Unds'in'Unds'keys'LParUndsRParUnds'MAP'Unds'Bool'Unds'KItem'Unds'Map"
-                []
-                [SortKItem, SortMap]
-                SortBool
-                ( Booster.Definition.Attributes.Base.SymbolAttributes
-                    Booster.Definition.Attributes.Base.TotalFunction
-                    Booster.Definition.Attributes.Base.IsNotIdem
-                    Booster.Definition.Attributes.Base.IsNotAssoc
-                    Booster.Definition.Attributes.Base.IsNotMacroOrAlias
-                    Booster.Definition.Attributes.Base.CanBeEvaluated
-                    Nothing
-                    Nothing
-                )
-            )
-            []
-            [k, m]
-    | otherwise = case sortOfTerm m of
-        SortVar{} -> error "maformed map sort"
-        SortApp nm _ ->
-            case stripPrefix "Sort" nm of
-                Nothing -> error "maformed map sort"
-                Just mapName -> case Map.lookup ("Lbl" <> mapName <> "'Coln'in'Unds'keys") symbols of
-                    Just inKeysSymbol ->
-                        SymbolApplication inKeysSymbol [] [k, m]
-                    Nothing -> error "in_keys for this map sort does not exist"
+mkInKeysMap :: Map.Map SymbolName Symbol -> Map.Map (Sort, Sort) Symbol
+mkInKeysMap symbols =
+    Map.fromList
+        [ (sorts, sym)
+        | sym@Symbol{argSorts, attributes = SymbolAttributes{hook}} <- Map.elems symbols
+        , hook == Just "MAP.in_keys"
+        , sorts <- mTuple argSorts
+        ]
+  where
+    mTuple [x, y] = [(x, y)]
+    mTuple _ = []
+
+mkInKeys :: Map.Map (Sort, Sort) Symbol -> Term -> Term -> Term
+mkInKeys inKeysSymbols k m =
+    case Map.lookup (sortOfTerm k, sortOfTerm m) inKeysSymbols of
+        Just inKeysSymbol -> SymbolApplication inKeysSymbol [] [k, m]
+        Nothing ->
+            error $
+                "in_keys for key sort '"
+                    <> show (pretty $ sortOfTerm k)
+                    <> "' and map sort '"
+                    <> show (pretty $ sortOfTerm m)
+                    <> "' does not exist."

library/Booster/Pattern/Base.hs

@@ -190,6 +190,7 @@ unitSymbol def =
                 , hasEvaluators = CanBeEvaluated
                 , collectionMetadata = Just def
                 , smt = Nothing
+                , hook = Nothing
                 }
         }
   where
@@ -213,6 +214,7 @@ concatSymbol def =
                 , hasEvaluators = CanBeEvaluated
                 , collectionMetadata = Just def
                 , smt = Nothing
+                , hook = Nothing
                 }
         }
   where
@@ -238,6 +240,7 @@ kmapElementSymbol def =
                 , hasEvaluators = CanBeEvaluated
                 , collectionMetadata = Just $ KMapMeta def
                 , smt = Nothing
+                , hook = Nothing
                 }
         }
 
@@ -257,6 +260,7 @@ klistElementSymbol def =
                 , hasEvaluators = CanBeEvaluated
                 , collectionMetadata = Just $ KListMeta def
                 , smt = Nothing
+                , hook = Nothing
                 }
         }
 
@@ -693,6 +697,7 @@ injectionSymbol =
                 , hasEvaluators = CanBeEvaluated
                 , collectionMetadata = Nothing
                 , smt = Nothing
+                , hook = Nothing
                 }
         }
 
@@ -723,7 +728,16 @@ pattern KSeq sort a =
                 []
                 [SortKItem, SortK]
                 SortK
-                (SymbolAttributes Constructor IsNotIdem IsNotAssoc IsNotMacroOrAlias CanBeEvaluated Nothing Nothing)
+                ( SymbolAttributes
+                        Constructor
+                        IsNotIdem
+                        IsNotAssoc
+                        IsNotMacroOrAlias
+                        CanBeEvaluated
+                        Nothing
+                        Nothing
+                        Nothing
+                    )
             )
         []
         [ Injection sort SortKItem a
@@ -733,7 +747,16 @@ pattern KSeq sort a =
                         []
                         []
                         SortK
-                        (SymbolAttributes Constructor IsNotIdem IsNotAssoc IsNotMacroOrAlias CanBeEvaluated Nothing Nothing)
+                        ( SymbolAttributes
+                                Constructor
+                                IsNotIdem
+                                IsNotAssoc
+                                IsNotMacroOrAlias
+                                CanBeEvaluated
+                                Nothing
+                                Nothing
+                                Nothing
+                            )
                     )
                 []
                 []

library/Booster/Pattern/Binary.hs

@@ -216,6 +216,7 @@ lookupKoreDefinitionSymbol name = DecodeM $ do
                         CannotBeEvaluated
                         Nothing
                         Nothing
+                        Nothing
                     )
         Just def -> Right $ Map.lookup name $ symbols def
 

library/Booster/Pattern/Bool.hs

@@ -61,6 +61,7 @@ pattern TotalFunctionWithSMT hook =
         CanBeEvaluated
         Nothing
         (Just (SMTHook (Atom (SMTId hook))))
+        Nothing
 
 pattern AndBool :: Term -> Term -> Term
 pattern AndBool l r =
@@ -148,6 +149,7 @@ pattern SetIn a b =
                         CanBeEvaluated
                         Nothing
                         Nothing
+                        Nothing
                     )
             )
         []

library/Booster/Syntax/Json/Internalise.hs

@@ -241,6 +241,7 @@ internaliseTermRaw qq allowAlias checkSubsorts sortVars definition@KoreDefinitio
                                     , hasEvaluators = Flag False
                                     , collectionMetadata = Nothing
                                     , smt = Nothing
+                                    , hook = Nothing
                                     }
                     else
                         maybe (throwE $ UnknownSymbol name symPatt) pure $