Merge dev into counter-example-generation

Author: kksnowy

ChangeLog.md

@@ -109,6 +109,8 @@ the ITP backend code can be invoked from any location.
 
 ### Loss backend
 
+* `--declaration` now accepts non-property declarations and restricts output to exactly the names listed.
+
 * Added the ability to declare custom Differentiable Logics internally in Vehicle (see documentation for details).
 
 * Fixed a bug where the compiler was erroring on some uses of `forall` for indices.

vehicle/src/Vehicle/Backend/Loss.hs

@@ -6,6 +6,8 @@ where
 import Control.Monad.Reader (ReaderT)
 import Data.Maybe (maybeToList)
 import Data.Proxy (Proxy (..))
+import Data.Set (Set)
+import Data.Set qualified as Set
 import Vehicle.Backend.Loss.Core
 import Vehicle.Backend.Loss.Domain (compileQuantifier)
 import Vehicle.Backend.Loss.LogicCompilation (findAndCompileLogic)
@@ -18,6 +20,7 @@ import Vehicle.Compile.Prelude
 import Vehicle.Data.Builtin.Loss
 import Vehicle.Data.Builtin.Standard
 import Vehicle.Data.Builtin.Standard.Normalise ()
+import Vehicle.Data.Code.BooleanExpr (unDisjunctAll)
 import Vehicle.Data.Code.ForcedValue
 import Vehicle.Data.DifferentiableLogic
 import Vehicle.Data.Variable.Bound.Context.Tensor (TensorBoundContextT)
@@ -26,49 +29,52 @@ import Vehicle.Data.Variable.Free.Context (MonadFreeContext (..), addDeclEntryTo
 convertToLossTensors ::
   (MonadCompile m) =>
   DifferentiableLogicID ->
+  Set Name ->
   Prog Builtin ->
   m (Prog LossBuiltin)
-convertToLossTensors logicID prog@(Main ds) = do
+convertToLossTensors logicID requestedDecls prog@(Main ds) = do
   -- First find and compile the logic
   logic <- logCompilerPass LossLogic $ findAndCompileLogic logicID prog
 
   -- Then compile the program using that logic
   runFreshFreeContextT (Proxy @Builtin) $ do
     runFreshFreeContextT (Proxy @LossBuiltin) $
       logCompilerPass Loss $ do
-        Main <$> convertDecls logicID logic ds
+        Main <$> convertDecls logicID logic requestedDecls ds
 
 convertDecls ::
   (MonadCompile m, MonadFreeContext Builtin m, MonadFreeContext LossBuiltin m) =>
   DifferentiableLogicID ->
   DifferentiableLogicImplementation ->
+  Set Name ->
   [Decl Builtin] ->
   m [Decl LossBuiltin]
-convertDecls logicID logic = \case
+convertDecls logicID logic requestedDecls = \case
   [] -> return []
   decl : decls -> do
-    maybeLossDecl <- convertDecl logicID logic decl
+    maybeLossDecl <- convertDecl logicID logic requestedDecls decl
     decls' <-
       maybe id addDeclToContext maybeLossDecl $
         addDeclEntryToContext decl $
-          convertDecls logicID logic decls
+          convertDecls logicID logic requestedDecls decls
     return $ maybeToList maybeLossDecl ++ decls'
 
 convertDecl ::
   forall m.
   (MonadCompile m, MonadFreeContext Builtin m, MonadFreeContext LossBuiltin m) =>
   DifferentiableLogicID ->
   DifferentiableLogicImplementation ->
+  Set Name ->
   Decl Builtin ->
   m (Maybe (Decl LossBuiltin))
-convertDecl logicID logic decl = case decl of
+convertDecl logicID logic requestedDecls decl = case decl of
   DefAbstract p ident sort typ
     | isAnnotatedAsExternalResource sort -> do
         let normType = Unforced emptyBoundEnv typ
         runConversion $ convertResourceDecl p ident sort normType
     | otherwise -> return Nothing
   DefFunction p ident ann typ expr
-    | isAnnotatedAsProperty ann -> do
+    | isAnnotatedAsProperty ann || nameOf decl `Set.member` requestedDecls -> do
         let normType = Unforced emptyBoundEnv typ
         let normExpr = Unforced emptyBoundEnv expr
         runConversion $ convertPropertyDecl p ident ann normType normExpr
@@ -104,11 +110,14 @@ convertPropertyDecl ::
 convertPropertyDecl p ident ann typ body = do
   lossType <- convertDeclType typ
   lossBody <- convertMultiProperty body
-  let lossTensorDecl = DefFunction p ident ann lossType lossBody
-  return lossTensorDecl
+  return $ DefFunction p ident ann lossType lossBody
 
 convertDeclType :: (MonadLogic m) => UnforcedType Builtin -> m (Type LossBuiltin)
 convertDeclType typ = unnormalise 0 <$> convertThunk Nothing typ
 
 convertMultiProperty :: (MonadLogic m) => Thunk Builtin -> m (Expr LossBuiltin)
-convertMultiProperty body = unnormalise 0 <$> convertThunk (Just compileQuantifier) body
+convertMultiProperty body = do
+  let compQuantifier args = do
+        disjuncts <- compileQuantifier args
+        foldrM1 orLossValue $ unDisjunctAll disjuncts
+  unnormalise 0 <$> convertThunk (Just compQuantifier) body

vehicle/src/Vehicle/Backend/Loss/Domain.hs

@@ -1,5 +1,6 @@
 module Vehicle.Backend.Loss.Domain
   ( compileQuantifier,
+    orLossValue,
   )
 where
 
@@ -12,7 +13,7 @@ import Data.List.NonEmpty (NonEmpty (..))
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Vehicle.Backend.Loss.Core
-import Vehicle.Backend.Loss.Domain.PurifyAssertion (BlockingReason, tryPurifyAssertion, unblockingActions)
+import Vehicle.Backend.Loss.Domain.PurifyAssertion (tryPurifyAssertion, unblockingActions)
 import Vehicle.Backend.Loss.LossCompilation
 import Vehicle.Backend.Solver.UserVariableElimination.ConstraintSearch (findAllBounds)
 import Vehicle.Compile.Constants.ForcedValue
@@ -50,17 +51,17 @@ import Vehicle.Prelude.Warning (CompileWarning (..))
 compileQuantifier ::
   (MonadLogic m) =>
   (Quantifier, QuantifyRatTensorArgs (Thunk Builtin) (Closure Builtin)) ->
-  m (Thunk LossBuiltin)
+  m (DisjunctAll (Thunk LossBuiltin))
 compileQuantifier (q, args) = do
   maybePartitions <- compileQuantifierInternal (q, args)
   case maybePartitions of
-    Trivial b ->
+    Trivial b -> do
       -- TODO add a warning
-      Forced <$> convertBoolTensorLiteral (ZeroDimTensor b)
+      value <- Forced <$> convertBoolTensorLiteral (ZeroDimTensor b)
+      return $ DisjunctAll [value]
     NonTrivial partitions -> do
       let disjunctedPartitions = partitionsToDisjuncts partitions
-      DisjunctAll (v :| vs) <- traverse checkFinalPartitionUnconstrained disjunctedPartitions
-      foldrM orLossValue v vs
+      traverse checkFinalPartitionUnconstrained disjunctedPartitions
 
 checkFinalPartitionUnconstrained ::
   (MonadLogic m) =>
@@ -292,21 +293,6 @@ type MonadDomain m =
     MonadTensorBoundContext m
   )
 
-orLossValue :: (MonadDomain m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
-orLossValue e1 e2 =
-  convertBooleanOp PointwiseDisjunction $
-    mkExpr accessSpine (TensorOp2Args (Forced IDimNil) e1 e2)
-
-andLossValue :: (MonadDomain m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
-andLossValue e1 e2 =
-  convertBooleanOp PointwiseConjunction $
-    mkExpr accessSpine (TensorOp2Args (Forced IDimNil) e1 e2)
-
-notLossValue :: (MonadDomain m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
-notLossValue dims e =
-  convertBooleanOp PointwiseNegation $
-    mkExpr accessSpine (TensorOp1Args dims e)
-
 notConstraint :: (MonadDomain m) => UserVariableConstraint LossBuiltin -> m (BooleanExpr (UserVariableConstraint LossBuiltin))
 notConstraint (NormalisedRelation rel expr) = do
   negExpr <- scaleExpr (-1) expr

vehicle/src/Vehicle/Backend/Loss/Domain/PurifyAssertion.hs

@@ -14,6 +14,7 @@ import Control.Applicative (liftA2)
 import Control.Monad (liftM2)
 import Control.Monad.Except (MonadError (..), runExceptT)
 import Vehicle.Compile.Constants.ForcedValue (TensorValueLinearExpr)
+import Vehicle.Compile.Error
 import Vehicle.Compile.Normalise.Force
 import Vehicle.Compile.Normalise.RewriteRules (forceAndRewriteTensor)
 import Vehicle.Compile.Normalise.TypedValue
@@ -281,10 +282,6 @@ tryAndUnblock dims expr = do
     Right unblocked -> forIfTreeM unblocked $ \unblockedExpr ->
       compileLinearExpr dims unblockedExpr
 
-data BlockingReason
-  = BlockingNetwork Identifier
-  | BlockingDatasetOrParameter Identifier
-
 unblockingActions ::
   (MonadPurifyAssertion m, MonadError BlockingReason m) =>
   UnblockingActions m
@@ -293,18 +290,19 @@ unblockingActions =
     { unblockRatTensorBoundVar = purifyBoundVar,
       unblockRecordBoundVar = purifyBoundVar,
       unblockNetworkApp = \_ _ ident _ -> throwError $ BlockingNetwork ident,
-      unblockDatasetOrParameter = \ident -> throwError $ BlockingDatasetOrParameter ident
+      unblockDatasetOrParameter = \_ ident -> throwError $ BlockingDatasetOrParameter ident
     }
 
 purifyBoundVar ::
-  (MonadLogger m, MonadReadableTensorBoundContext m) =>
+  (MonadPurifyAssertion m) =>
+  TypeUnblockingFunction (Thunk Builtin) m ->
   Lv ->
-  m (Thunk Builtin)
-purifyBoundVar lv = do
+  m (IfTree (Thunk Builtin) (Thunk Builtin))
+purifyBoundVar unblock lv = do
   (_, maybeChildVars) <- lookupVariableInNestedCtx lv
   case maybeChildVars of
-    Nothing -> return $ Forced $ VBoundVar lv []
-    Just (_tensorVar, sliceVar) -> replaceTensorVariableWithStackedChildren sliceVar
+    Nothing -> return $ IfLeaf $ Forced $ VBoundVar lv []
+    Just (_tensorVar, sliceVar) -> unblock =<< replaceTensorVariableWithStackedChildren sliceVar
 
 --------------------------------------------------------------------------------
 -- Utility functions

vehicle/src/Vehicle/Backend/Loss/JSON.hs

@@ -7,6 +7,7 @@ module Vehicle.Backend.Loss.JSON
   )
 where
 
+import Control.Monad.Except (MonadError (..))
 import Data.Aeson (ToJSON (..), genericToJSON)
 import Data.List (elemIndex)
 import Data.Map (Map)
@@ -35,7 +36,7 @@ import Vehicle.Data.Code.Interface.Args
 import Vehicle.Data.Tensor (ExtendedRatTensor)
 import Vehicle.Data.Variable.Bound.Context.Name
 import Vehicle.Data.Variable.Free.Context (MonadFreeContext, addDeclToContext, runFreshFreeContextT)
-import Vehicle.Prelude (Doc, GenericArg (..), HasName (..), HasType (..), Identifier (..), Name, Provenance, explicit, indent, jsonOptions, line, mkExplicitBinder, resolutionError, squotes, userModulePath)
+import Vehicle.Prelude (Doc, GenericArg (..), HasName (..), HasType (..), Identifier (..), Name, Provenance, explicit, indent, jsonOptions, line, mkExplicitBinder, resolutionError, squotes, stdlibIdentifier, userModulePath)
 import Vehicle.Prelude.Error (developerError)
 import Vehicle.Prelude.Logging.Class
 
@@ -166,8 +167,8 @@ type MonadJSON m =
     MonadFreeContext LossBuiltin m
   )
 
-unsupportedError :: (Pretty a) => a -> b
-unsupportedError b = developerError $ "Conversion of" <+> pretty b <+> "is not yet implemented"
+unsupportedError :: (MonadJSON m, Pretty a) => a -> m b
+unsupportedError a = throwError $ UnsupportedLossOperation (stdlibIdentifier "unknown", mempty) (pretty a)
 
 dependentTypesError :: (Pretty a) => a -> b
 dependentTypesError b = developerError $ "Conversion of" <+> pretty b <+> "is not yet implemented"

vehicle/src/Vehicle/Backend/Loss/LogicCompilation.hs

@@ -6,7 +6,7 @@ module Vehicle.Backend.Loss.LogicCompilation
 where
 
 import Control.Monad (foldM)
-import Control.Monad.Except (MonadError (..))
+import Control.Monad.Except (MonadError (..), runExceptT)
 import Control.Monad.State (MonadState, StateT, execStateT, modify)
 import Data.Map (Map)
 import Data.Map qualified as Map
@@ -17,12 +17,11 @@ import Vehicle.Backend.Loss.Core hiding (lookupLogicField)
 import Vehicle.Backend.Loss.LossCompilation (convertQuantifierlessExprToLoss)
 import Vehicle.Backend.Prelude (DifferentiableLogicID)
 import Vehicle.Compile.Error
-import Vehicle.Compile.Normalise.Builtin
-import Vehicle.Compile.Normalise.Core
 import Vehicle.Compile.Normalise.Force
 import Vehicle.Compile.Normalise.Quote (unnormalise)
 import Vehicle.Compile.Prelude
 import Vehicle.Compile.Print
+import Vehicle.Compile.Unblock (noUnblocking, unblockBoolExpr)
 import Vehicle.Data.Builtin.Interface (Accessor (..))
 import Vehicle.Data.Builtin.Loss (ComparisonOp (..), LogicDirection, LossBuiltin)
 import Vehicle.Data.Builtin.Standard (Builtin)
@@ -153,13 +152,17 @@ calculateLogicDirection ::
   OMap FieldName (Thunk Builtin) ->
   m LogicDirection
 calculateLogicDirection declProv fields = do
-  let trueValue = lookupLogicField TruthityElement fields
-  let falseValue = lookupLogicField FalsityElement fields
-  let args = TensorComparisonArgs (Forced IDimNil) (Forced IDimNil) trueValue falseValue
-  result <- runFreshNameBoundContextT $ evalCompareRatTensor @ForcedValue Le args
-  case result of
-    Evaluated (Forced (IBoolLiteral b)) -> return b
-    _ -> throwError $ UnorderableDifferentiableLogic declProv (mkExpr accessCompareRatTensor (Le, args))
+  let expr = do
+        let trueValue = lookupLogicField TruthityElement fields
+        let falseValue = lookupLogicField FalsityElement fields
+        let args = TensorComparisonArgs (Forced IDimNil) (Forced IDimNil) trueValue falseValue
+        Forced $ mkExpr accessCompareRatTensor (Le, args)
+
+  errorOrResult <- runExceptT $ runFreshNameBoundContextT $ forceThunk =<< unblockBoolExpr noUnblocking expr
+  case errorOrResult of
+    Left blockingErr -> throwError $ UnorderableDifferentiableLogic declProv expr (Left blockingErr)
+    Right (IBoolLiteral b) -> return b
+    Right result -> throwError $ UnorderableDifferentiableLogic declProv expr (Right result)
 
 compileLogicField ::
   (MonadLoss m) =>

vehicle/src/Vehicle/Backend/Loss/LossCompilation.hs

@@ -3,6 +3,9 @@ module Vehicle.Backend.Loss.LossCompilation
     convertBoolTensorLiteral,
     convertQuantifierlessExprToLoss,
     convertBooleanOp,
+    orLossValue,
+    andLossValue,
+    notLossValue,
   )
 where
 
@@ -269,6 +272,24 @@ convertBoolTensorLiteral tensor = do
 --------------------------------------------------------------------------------
 -- Utils
 
+orLossValue :: (MonadLogic m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
+orLossValue e1 e2 =
+  convertBooleanOp PointwiseDisjunction $
+    mkExpr accessSpine (TensorOp2Args (Forced IDimNil) e1 e2)
+
+andLossValue :: (MonadLogic m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
+andLossValue e1 e2 =
+  convertBooleanOp PointwiseConjunction $
+    mkExpr accessSpine (TensorOp2Args (Forced IDimNil) e1 e2)
+
+notLossValue :: (MonadLogic m) => Thunk LossBuiltin -> Thunk LossBuiltin -> m (Thunk LossBuiltin)
+notLossValue dims e =
+  convertBooleanOp PointwiseNegation $
+    mkExpr accessSpine (TensorOp1Args dims e)
+
+--------------------------------------------------------------------------------
+-- Utils
+
 currentPass :: Doc a
 currentPass = "logic translation"
 

vehicle/src/Vehicle/Backend/Solver.hs

@@ -30,6 +30,7 @@ import Vehicle.Compile.Print (prettyFriendly, prettyFriendlyEmptyCtx)
 import Vehicle.Compile.Print.Warning ()
 import Vehicle.Compile.Property (traverseMultiProperty)
 import Vehicle.Compile.Unblock (UnblockingActions (..), unblockBoolExpr)
+import Vehicle.Data.Builtin.Interface (Accessor (..))
 import Vehicle.Data.Builtin.Standard
 import Vehicle.Data.Code.BooleanExpr
 import Vehicle.Data.Code.ForcedValue
@@ -317,10 +318,10 @@ compileQuerySetPartitions globalCtx isPropertyNegated maybePartitions = case may
 topLevelUnblockingActions :: (Monad m) => UnblockingActions m
 topLevelUnblockingActions =
   UnblockingActions
-    { unblockRatTensorBoundVar = developerError "No bound variables should exist at top-level",
-      unblockRecordBoundVar = developerError "No bound variables should exist at top-level",
-      unblockNetworkApp = \_ _ _ -> developerError "Unblocking of constant network functions at top-level not yet supported",
-      unblockDatasetOrParameter = developerError "Should not be unblocking datasets or parameters"
+    { unblockRatTensorBoundVar = \_ _ -> developerError "No bound variables should exist at top-level",
+      unblockRecordBoundVar = \_ _ -> developerError "No bound variables should exist at top-level",
+      unblockNetworkApp = \_ _ ident args -> return $ IfLeaf $ Forced $ VFreeVar ident (mkExpr accessSpine args),
+      unblockDatasetOrParameter = \_ _ -> developerError "Should not be unblocking datasets or parameters"
     }
 
 handlePropertyCompileError ::

vehicle/src/Vehicle/Backend/Solver/UserVariableElimination.hs

@@ -38,7 +38,7 @@ import Vehicle.Compile.Unblock qualified as Unblocking
 import Vehicle.Compile.Variable (createUserVar)
 import Vehicle.Data.Builtin.Interface
 import Vehicle.Data.Builtin.Standard
-import Vehicle.Data.Code.BooleanExpr (elimIfTree)
+import Vehicle.Data.Code.BooleanExpr (IfTree, elimIfTree)
 import Vehicle.Data.Code.ForcedValue
 import Vehicle.Data.Code.Interface
 import Vehicle.Data.MaybeTrivial
@@ -271,15 +271,16 @@ unblockingActions =
     { unblockRatTensorBoundVar = unblockQuantifiedBoundVar,
       unblockRecordBoundVar = unblockQuantifiedBoundVar,
       unblockNetworkApp = unblockNetworkApplication,
-      unblockDatasetOrParameter = unexpectedExprError "solver compilation" "dataset or parameter"
+      unblockDatasetOrParameter = \_ _ -> unexpectedExprError "solver compilation" "dataset or parameter"
     }
 
 unblockQuantifiedBoundVar ::
   (MonadQuantifierBody m) =>
+  TypeUnblockingFunction (Thunk Builtin) m ->
   Lv ->
-  m (Thunk Builtin)
-unblockQuantifiedBoundVar lv =
-  replaceTensorVariableWithStackedChildren (SliceVariable lv)
+  m (IfTree (Thunk Builtin) (Thunk Builtin))
+unblockQuantifiedBoundVar unblock lv =
+  unblock =<< replaceTensorVariableWithStackedChildren (SliceVariable lv)
 
 unblockNetworkApplication ::
   (MonadQuantifierBody m) =>

vehicle/src/Vehicle/Backend/Solver/UserVariableElimination/PurifyAssertion.hs

@@ -104,7 +104,7 @@ purifyBoundVar ::
   Lv ->
   PurifyFn m
 purifyBoundVar v actions@UnblockingActions {..} incrDims
-  | incrDims > 0 = purifyExpr actions incrDims =<< unblockRatTensorBoundVar v
+  | incrDims > 0 = unblockRatTensorBoundVar (purifyExpr actions incrDims) v
   | otherwise = return $ IfLeaf $ Forced $ VBoundVar v []
 
 purifyNetworkVar ::