diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2023, Orestis Melkonian +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,10 @@ +agda2train: An Agda backend to generate training data for machine learning +=========================================================================== + +This is a work in progress. + +## Relevant Agda issues + +- [Type-on-hover #516](https://github.com/agda/agda/issues/516) +- [Step-by-step evaluation #4747](https://github.com/agda/agda/issues/4747) +- [Reverting projection-like optimization #5142](https://github.com/agda/agda/pull/5142/) diff --git a/agda2train.cabal b/agda2train.cabal new file mode 100644 --- /dev/null +++ b/agda2train.cabal @@ -0,0 +1,62 @@ +cabal-version: 2.4 +name: agda2train +version: 0.0.1.0 +synopsis: Agda backend to generate training data for machine learning purposes. +description: - +category: Agda, machine learning +author: Orestis Melkonian +maintainer: melkon.or@gmail.com +stability: experimental +homepage: http://github.com/omelkonian/agda2train/ +bug-reports: https://github.com/omelkonian/agda2train/issues +license: BSD-3-Clause +license-file: LICENSE +copyright: (c) 2023 Orestis Melkonian + +extra-source-files: + data/defsToSkip.txt + README.md + LICENSE + +source-repository head + type: git + location: git://github.com/omelkonian/agda2train.git + + + +executable agda2train + main-is: Main.hs + other-modules: ToTrain, Output, AgdaInternals + build-depends: base >=4.12.0.0 && <4.20 + , Agda >=2.6.3 && <2.6.4 + , pretty >=1.1.3.3 && <1.2 + , containers >=0.5.11.0 && <0.8 + , unordered-containers >=0.2.9.0 && <0.3 + , aeson >=2.0 && <2.3 + , aeson-pretty == 0.8.10 + , bytestring >=0.10.8.1 && <0.13 + , mtl >=2.2.1 && <2.4 + , directory >=1.2.6.2 && <1.4 + , filepath >=1.4.1.0 && <1.5 + , deepseq >=1.4.2.0 && <1.6 + , async >=2.2 && <2.3 + , file-embed == 0.0.15.0 + hs-source-dirs: src + default-language: Haskell2010 + default-extensions: LambdaCase + , NamedFieldPuns + , OverloadedStrings + , RecordWildCards + , TupleSections + , DeriveGeneric + , DeriveAnyClass + , DeriveFunctor + , DeriveFoldable + , DeriveTraversable + , DerivingVia + , StandaloneDeriving + , PatternSynonyms + , TypeFamilies + , TypeOperators + , FunctionalDependencies + ghc-options: -threaded -rtsopts -with-rtsopts=-N diff --git a/data/defsToSkip.txt b/data/defsToSkip.txt new file mode 100644 --- /dev/null +++ b/data/defsToSkip.txt @@ -0,0 +1,47 @@ +Data.Relation.Binary.Subset.Propositional.Properties.filter⁺ +foundation.homotopies._.compute-ind-htpy +foundation.equivalences.coh-unit-laws-equiv +foundation.functoriality-propositional-truncation.id-map-trunc-Prop +foundation.functoriality-coproduct-types._.equiv-mutually-exclusive-coprod +trees.polynomial-endofunctors.coh-refl-htpy-polynomial-endofunctor +foundation.uniqueness-image._.htpy-map-hom-equiv-slice-uniqueness-im +foundation.uniqueness-image._.tetrahedron-hom-equiv-slice-uniqueness-im +foundation.functoriality-set-truncation._.htpy-map-hom-slice-trunc-im-Set +foundation.functoriality-set-truncation._.tetrahedron-map-hom-slice-trunc-im-Set +univalent-combinatorics.2-element-types._.is-not-identity-swap-2-Element-Type +foundation.unordered-pairs._.eq-Eq-refl-unordered-pair +MGS.Equivalence-Induction.𝔾-≃-equation +MGS.More-FunExt-Consequences.EM-is-subsingleton +foundation.unordered-pairs.preserves-refl-htpy-unordered-pair +MGS.Equivalence-Constructions.Eq-Eq-cong +foundation.unordered-pairs.id-equiv-standard-unordered-pair +MGS.Embeddings.Embβ†’fun +trees.morphisms-algebras-polynomial-endofunctors._.refl-htpy-hom-algebra-polynomial-endofunctor +MGS.Universe-Lifting._.q +UF.Equiv-FunExt.prop-＝-≃-⇔ +trees.morphisms-algebras-polynomial-endofunctors._.is-contr-total-htpy-hom-algebra-polynomial-endofunctor +trees.w-types.htpy-hom-π•Ž-Alg +UF.FunExt-Properties.naive-funext-gives-funextβ‚€ +graph-theory.morphisms-undirected-graphs._.refl-htpy-hom-Undirected-Graph +UF.UA-FunExt._._.h +graph-theory.morphisms-undirected-graphs._.is-contr-total-htpy-hom-Undirected-Graph +UF.PairFun.pair-fun-embedding +graph-theory.equivalences-undirected-graphs._.edge-standard-unordered-pair-vertices-id-equiv-Undirected-Graph +UF.UniverseEmbedding.hSet-embeddings-are-embeddings +graph-theory.equivalences-undirected-graphs._.is-contr-total-htpy-equiv-Undirected-Graph +UF.Size._.Set-Replacement +Lifting.IdentityViaSIP.𝓛-Id +foundation.partitions._.is-section-map-inv-compute-block-partition +Lifting.IdentityViaSIP.𝓛-IdΒ· +foundation.partitions._.is-retraction-map-inv-compute-block-partition +Lifting.Size._.e +foundation.partitions._.is-emb-inhabited-subtype-block-partition +InjectiveTypes.Blackboard.injective._.e +foundation.partitions._.compute-is-in-block-partition +foundation.partitions._.is-contr-block-containing-element-partition +foundation.partitions._.is-in-block-class-partition +foundation.partitions._.compute-total-block-partition +TypeTopology.SquashedCantor._.fg +foundation.partitions.partition-Set-Indexed-Ξ£-Decomposition +foundation.set-quotients._.is-section-equivalence-class-set-quotient +foundation.set-quotients._.is-retraction-equivalence-class-set-quotient diff --git a/src/AgdaInternals.hs b/src/AgdaInternals.hs new file mode 100644 --- /dev/null +++ b/src/AgdaInternals.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-} +-- | Missing instances for WriterT in Agda's internals. +-- +-- NB: we are not exporting any definitions, just typeclass instances. +module AgdaInternals () where + +import Data.String + +import Control.Monad +import Control.Monad.Writer +import Control.Monad.State + +import Agda.TypeChecking.Monad +import Agda.TypeChecking.Warnings +import Agda.TypeChecking.Pretty + +instance (Monoid r, Monad m, Semigroup (m Doc)) => Semigroup (WriterT r m Doc) where + (<>) = liftM2 (<>) +instance (Monoid r, Monad m, IsString (m a)) => IsString (WriterT r m a) where + fromString s = WriterT $ fmap (,mempty) $ fromString s +instance (Monoid r, MonadStConcreteNames m) => MonadStConcreteNames (WriterT r m) where + runStConcreteNames m = WriterT $ runStConcreteNames $ StateT $ \s -> do + ((x,s'),ns) <- runWriterT $ runStateT m s + return ((x,ns),s') +instance (Monoid r, MonadInteractionPoints m) => MonadInteractionPoints (WriterT r m) +instance (Monoid r, MonadFresh i m) => MonadFresh i (WriterT r m) +instance (Monoid r, MonadWarning m) => MonadWarning (WriterT r m) where +instance (Monoid r, MonadBlock m) => MonadBlock (WriterT r m) where + catchPatternErr h m = WriterT $ catchPatternErr (runWriterT . h) (runWriterT m) +instance (Monoid r, MonadTCM m, MonadConstraint m) => MonadConstraint (WriterT r m) where + addConstraint x = lift . addConstraint x + addAwakeConstraint x = lift . addAwakeConstraint x + solveConstraint = lift . solveConstraint + solveSomeAwakeConstraints x = lift . solveSomeAwakeConstraints x + wakeConstraints = lift . wakeConstraints + stealConstraints = lift . stealConstraints + modifyAwakeConstraints = lift . modifyAwakeConstraints + modifySleepingConstraints = lift . modifySleepingConstraints +instance (Monoid r, MonadTCM m, MonadMetaSolver m) => MonadMetaSolver (WriterT r m) where + newMeta' x y z k l = lift . newMeta' x y z k l + assignV x y z w = lift . assignV x y z w + assignTerm' x y = lift . assignTerm' x y + etaExpandMeta x = lift . etaExpandMeta x + updateMetaVar x = lift . updateMetaVar x + speculateMetas x = WriterT . runWriterT . speculateMetas x diff --git a/src/Main.hs b/src/Main.hs new file mode 100644 --- /dev/null +++ b/src/Main.hs @@ -0,0 +1,234 @@ +{-# LANGUAGE TypeApplications #-} +-- | Main driver for the @agda2train@ executable. +module Main where + +import GHC.Generics + +import System.Environment ( getArgs, withArgs ) +import System.Directory ( createDirectoryIfMissing, doesFileExist ) +import System.FilePath ( () ) +import System.Console.GetOpt ( OptDescr(..), ArgDescr(..) ) + +import qualified Data.List as L +import qualified Data.Set as S +import qualified Data.HashMap.Strict as HM +import qualified Data.ByteString.Lazy as BL +import Data.Aeson ( ToJSON ) +import Data.Aeson.Encode.Pretty + ( encodePretty', Config(..), defConfig, Indent(..), keyOrder ) + +import Control.Monad +import Control.Monad.IO.Class ( liftIO ) +import Control.DeepSeq ( NFData ) + +import Agda.Main ( runAgda ) +import Agda.Compiler.Backend hiding (Reduced) +import Agda.Compiler.Common ( curDefs ) + +import Agda.Utils.Pretty +import Agda.Utils.Maybe +import Agda.Utils.Monad + +import Agda.Syntax.Scope.Base +import Agda.Syntax.Scope.Monad +import Agda.Syntax.TopLevelModuleName + +import Agda.TypeChecking.Reduce +import qualified Agda.TypeChecking.Pretty as P + +import Agda.Compiler.Common ( curIF ) + +import ToTrain +import Output + +-- | The main entrypoint for the @agda2train@ executable. +main :: IO () +main = do + as <- getArgs + let extraFlags = [] -- ["--no-projection-like"] + withArgs (extraFlags ++ as) $ + runAgda [Backend $ mkBackend train] + +-- | Make an Agda backend from a given training function (c.f. 'ToTrain.TrainF'). +mkBackend :: TrainF -> + Backend' Options Options ([ScopeEntry], [ScopeEntry]) () (String, [Sample]) +mkBackend trainF = Backend' + { backendName = "agda2train" + , backendVersion = Nothing + , options = defaultOptions + , commandLineFlags = + [ Option ['r'] ["recurse"] (NoArg recOpt) + "Recurse into imports/dependencies." + , Option ['x'] ["no-json"] (NoArg noJsonOpt) + "Skip generation of JSON files. (just debug print)" + , Option [] ["ignore-existing-json"] (NoArg ignoreJsonOpt) + "Ignore existing JSON files. (i.e. always overwrite)" + , Option [] ["print-json"] (NoArg printJsonOpt) + "Print JSON output. (for debugging)" + , Option [] ["no-terms"] (NoArg includeDefsOpt) + "Do not include definitions of things in scope" + , Option [] ["no-privates"] (NoArg includePrivsOpt) + "Do not include private definitions" + , Option ['o'] ["out-dir"] (ReqArg outDirOpt "DIR") + "Generate data at DIR. (default: project root)" + ] + , isEnabled = \ _ -> True + , scopeCheckingSuffices = False + , mayEraseType = \ _ -> return True + , preCompile = return + , postCompile = \ _ _ _ -> return () + , preModule = \ opts isMain md _ -> + let + processScopeEntry :: QName -> TCM (Maybe ScopeEntry) + processScopeEntry qn = + caseMaybeM (tryMaybe $ do + def <- getConstInfo qn + case theDef def of + GeneralizableVar -> fail "not handling `variable` definitions" + DataOrRecSig _ -> fail "not handling `DataOrRecSig` definitions" + _ -> typeOfConst qn) (pure Nothing) $ \ty -> do + rty <- mkReduced ty + pty <- ppm ty + rty' <- traverse convert rty + + def <- getConstInfo qn + pdef <- ppm def + def' <- convert def + + report 20 $ ppm (pp qn) <> " : " <> ppm ty + report 30 $ " *pretty: " <> P.text (pp ty) + report 20 $ ppm (pp qn) <> " = " <> ppm def + report 30 $ " *pretty: " <> P.text (pp def) + reportReduced rty + report 20 "" + + return $ Just $ ppName qn :~ ScopeEntry + { _type = prender pty :> rty' + , definition = boolToMaybe (includeDefs opts) + $ prender pdef :> def' + , holes = Nothing + } + in + ifM (skip opts isMain md) (return $ Skip ()) $ do + report 10 $ "************************ " + <> ppm md <> " (" <> ppm (show isMain) + <> ") ***********************************" + -- printScope "" 1 "" + setScope . iInsideScope =<< curIF + -- ** public + NameSpace _ _ ns <- allThingsInScope <$> getCurrentScope + scopeEntries <- mapMaybeM processScopeEntry (S.toList ns) + -- ** private + privs <- HM.filterWithKey (\k _ -> k `S.notMember` ns) <$> liftTCM curDefs + privScopeEntries <- mapMaybeM processScopeEntry (HM.keys privs) + report 10 "******************************************************************" + return $ Recompile (scopeEntries, privScopeEntries) + , compileDef = \ _ _ _ def -> (ppName (defName def) ,) <$> runC (forEachHole trainF def) + , postModule = \ opts (scopeEntries, privScopeEntries) isMain md samples -> + let mn = pp md in + unlessM (skip opts isMain md) $ do + whenJust (outDir opts) (liftIO . createDirectoryIfMissing True) + unless (noJson opts) $ liftIO $ do + let + isGlobal = \(n :~ _) -> isNothing $ lookup n samples + (globals, locals') = L.partition isGlobal scopeEntries + locals = flip map locals' $ \(n :~ l) -> + n :~ l { holes = Just $ fromJust $ lookup n samples } + plocals = flip map privScopeEntries $ \(n :~ l) -> + n :~ l { holes = Just $ fromJust $ lookup n samples } + fileData = mn :~ TrainData + { scopeGlobal = globals + , scopeLocal = locals + , scopePrivate = boolToMaybe (includePrivs opts) plocals + } + + encodeFile (getOutFn opts mn) fileData + when (printJson opts) $ do + putStrLn "************************ JSON ************************" + BL.putStr (encode fileData) + putStrLn "******************************************************" + } + where + skip :: Options -> IsMain -> TopLevelModuleName -> TCM Bool + skip opts@Options{..} isMain md = do + jsonExists <- liftIO $ doesFileExist (getOutFn opts $ pp md) + return $ (not recurse && (isMain == NotMain)) + || (not noJson && jsonExists && not ignoreJson) + +-- * Command-line flags + +-- | Available options to configure how @agda2train@ generates training data: +-- +-- ['recurse'] whether to recursively generate data from all transitive dependencies +-- +-- ['noJson'] mock run without generating any actual JSON files +-- +-- ['ignoreJson'] compile everything from scratch +-- +-- ['includeDefs'] whether to include definitions in scope, or just their type +-- +-- ['includePrivs'] whether to include private definitions as well +-- +-- (run @agda2train -h/--help@ for a human-readable description of all options) +data Options = Options + { recurse, noJson, ignoreJson, printJson, includeDefs, includePrivs :: Bool + , outDir :: Maybe FilePath + } deriving (Generic, NFData) + +-- | The default options. +defaultOptions = Options + { recurse = False + , noJson = False + , ignoreJson = False + , printJson = False + , includeDefs = True + , includePrivs = True + , outDir = Nothing + } + +recOpt, noJsonOpt, ignoreJsonOpt, printJsonOpt, includeDefsOpt, includePrivsOpt + :: Monad m => Options -> m Options +recOpt opts = return $ opts { recurse = True } +noJsonOpt opts = return $ opts { noJson = True } +printJsonOpt opts = return $ opts { printJson = True } +ignoreJsonOpt opts = return $ opts { ignoreJson = True } +includeDefsOpt opts = return $ opts { includeDefs = False } +includePrivsOpt opts = return $ opts { includePrivs = False } + +outDirOpt :: Monad m => FilePath -> Options -> m Options +outDirOpt fp opts = return $ opts { outDir = Just fp } + +getOutDir :: Options -> FilePath +getOutDir opts = case outDir opts of + Just fp -> fp + Nothing -> "." + +getOutFn :: Options -> String -> FilePath +getOutFn opts mn = getOutDir opts mn <> ".json" + +-- * JSON encoding + +-- | Uses "Aeson.Pretty" to order the JSON fields. +encode :: ToJSON a => a -> BL.ByteString +encode = encodePretty' $ defConfig + { confIndent = Spaces 2 + , confCompare = keyOrder + [ "pretty" + , "tag" + , "name" + , "original", "simplified", "reduced", "normalised" + , "telescope", "patterns", "fields" + , "domain", "codomain" + , "abstraction", "body" + , "sort", "level", "literal" + , "head", "arguments" + , "variants", "reference", "variant" + , "index" + , "scopeGlobal", "scopeLocal" + , "type", "definition", "holes" + , "ctx", "goal", "term", "premises" + ] + } + +encodeFile :: ToJSON a => FilePath -> a -> IO () +encodeFile = \fn -> BL.writeFile fn . encode diff --git a/src/Output.hs b/src/Output.hs new file mode 100644 --- /dev/null +++ b/src/Output.hs @@ -0,0 +1,429 @@ +{-# LANGUAGE FlexibleInstances, FlexibleContexts #-} +-- | Defines the structure of the training data, as well as how to translate +-- internal Agda definition to this format. +module Output where + +import Control.Arrow ( second ) +import Control.Applicative ( (<|>) ) +import GHC.Generics ( Generic ) +import Data.List ( notElem, elemIndex ) +import Data.String ( fromString ) +import Data.Aeson +import qualified Data.Aeson as JSON +import qualified Data.Aeson.KeyMap as KM + +import Agda.Syntax.Common ( unArg ) +import qualified Agda.Syntax.Common as A +import qualified Agda.Syntax.Internal as A +import qualified Agda.Syntax.Literal as A +import Agda.Syntax.Internal + ( QName, absName, qnameName, qnameModule, unAbs, unEl, unDom + , nameId, conName, dbPatVarIndex, pDom, telToList, telFromList ) +import Agda.Syntax.Translation.InternalToAbstract ( NamedClause(..) ) +import qualified Agda.TypeChecking.Monad as A +import Agda.TypeChecking.Monad + ( TCM, MonadTCM, liftTCM, typeOfConst, theDef, defName, getConstInfo + , reportSDoc, VerboseLevel ) + +import qualified Agda.Utils.Pretty as P +import Agda.TypeChecking.Pretty + ( PrettyTCM(..), MonadPretty, fsep, punctuate, braces, parens, Doc ) +import qualified Agda.TypeChecking.Pretty as P + hiding (text) + +-- * Types + +-- | Name identifiers. +type Name = String +-- | DeBruijn indices. +type DB = Int +-- | A head of a Ξ»-application can either be a defined name in the global scope, +-- or a DeBruijn index into the local context. +type Head = Either Name DB + +-- * Generic constructions + +infixr 4 :>; pattern x :> y = Pretty {pretty = x, thing = y} +-- | Bundle a thing with its "pretty" version. +-- +-- NB: In JSON format, we follow a /shallow/ encoding with "pretty" being +-- an additional field and "thing" just inlined in the top-level record. +data Pretty a = Pretty + { pretty :: String + , thing :: a + } deriving Generic +deriving instance Show a => Show (Pretty a) +instance ToJSON a => ToJSON (Pretty a) where + toJSON (Pretty{..}) = let pretty' = toJSON pretty in + case toJSON thing of + (Object fs) -> object ("pretty" .= pretty' : KM.toList fs) + t@(Array xs) -> object ["pretty" .= pretty', "telescope" .= t] + t -> object ["pretty" .= pretty', "thing" .= toJSON t] +instance FromJSON a => FromJSON (Pretty a) where + parseJSON = withObject "Pretty" $ \v -> Pretty + <$> v .: "pretty" + <*> (v .: "thing" <|> parseJSON (Object v)) + +-- | Bundle a term with (several of) its normalised forms. +-- +-- We do not repeat identical elements in subsequent evaluations +-- (in the order simp/red/norm) and some fields may not be populated due to +-- the evaluation taking to long and leading to a timeout (currently 2 seconds). +-- +-- NB: Like 'Named', encoded in a /shallow/ JSON. +data Reduced a = Reduced + { original :: a + , simplified :: Maybe a + , reduced :: Maybe a + , normalised :: Maybe a + } deriving (Generic, Functor, Foldable, Traversable) +deriving instance Show a => Show (Reduced a) +instance ToJSON a => ToJSON (Reduced a) where + toJSON r@(Reduced{..}) + | Nothing <- simplified <|> reduced <|> normalised + = toJSON original + | otherwise + = genericToJSON jsonOpts r +instance FromJSON a => FromJSON (Reduced a) where + parseJSON = withObject "Object" $ \v -> Reduced + <$> (v .: "original" <|> parseJSON (Object v)) + <*> v .:? "simplified" + <*> v .:? "reduced" + <*> v .:? "normalised" + +infixr 4 :~; pattern x :~ y = Named {name = x, item = y} + +-- | Bundle a thing with its name. +-- +-- NB: Like 'Named' and 'Reduced', encoded in a /shallow/ JSON. +data Named a = Named + { name :: Name + , item :: a + } deriving (Generic, Show) +instance ToJSON a => ToJSON (Named a) where + toJSON (Named{..}) = let name' = toJSON name in + case toJSON item of + (Object fs) -> object ("name" .= name' : KM.toList fs) + x -> object ["name" .= name', "item" .= toJSON x] +instance FromJSON a => FromJSON (Named a) where + parseJSON = withObject "Named" $ \v -> Named + <$> v .: "name" + <*> (v .: "item" <|> parseJSON (Object v)) + +-- * Concrete types (~ JSON schema) + +-- | Data for a file include the filename and its training data. +type FileData = Named TrainData +-- | The training data for a module, divided into three parts. +data TrainData = TrainData + { scopeGlobal :: [ScopeEntry] + -- ^ The /global/ scope, giving the types and definitions of all @import@ statements. + -- + -- NB: does not contain any /holes/ for subterms. + , scopeLocal :: [ScopeEntry] + -- ^ The /local/ scope, containing the types, definitions, and training data + -- for each of this module's definitions. + , scopePrivate :: Maybe [ScopeEntry] + -- ^ The /private/ scope, containing private definitions not exported to the public, + -- as well as system-generated definitions stemming from @where@ or @with@. + } deriving Generic +instance ToJSON TrainData where toJSON = genericToJSON jsonOpts +instance FromJSON TrainData where parseJSON = genericParseJSON jsonOpts + +-- | Every 'ScopeEntry'' is /named/. +type ScopeEntry = Named ScopeEntry' +-- | An entry in the scope: type, definitions, and holes. +data ScopeEntry' = ScopeEntry + { _type :: Pretty (Reduced Type) + -- ^ The entry's type. + , definition :: Maybe (Pretty Definition) + -- ^ The actual body of this entry's definition. + , holes :: Maybe [Sample] + -- ^ Training data for each of the subterms in this entry's 'definition'. + } deriving (Generic, Show) +instance ToJSON ScopeEntry' where toJSON = genericToJSON jsonOpts +instance FromJSON ScopeEntry' where parseJSON = genericParseJSON jsonOpts + +-- | The training sample for each sub-hole. +data Sample = Sample + { ctx :: Pretty Telescope + -- ^ The current context, as a /binding telescope/. + , goal :: Pretty (Reduced Type) + -- ^ The current goal, i.e. type of the sub-term. + -- + -- NB: DeBruijn indices here refer to the 'ctx'. + , term :: Pretty (Reduced Term) + -- ^ The term that successfully fills the current 'goal'. + , premises :: [Name] + -- ^ Definitions used in this "proof", intended to be used for /premise selection/. + } deriving (Generic, Show, ToJSON, FromJSON) + + +-- | Agda definitions: datatypes, records, functions, postulates and primitives. +data Definition + = ADT {variants :: [Type]} + -- ^ e.g. + -- data β„• : Set where + -- zero : β„• + -- suc : β„• β†’ β„• + | Constructor {reference :: Name, variant :: Integer} + -- ^ e.g. `(β„•, 0) ~ zero` or `(β„•, 1) ~ suc` + | Record {telescope :: Telescope, fields :: [Type]} + -- ^ e.g. + -- record X : Set where + -- field x : β„• + -- y : β„• + | Function {clauses :: [Clause]} + -- ^ e.g. + -- f [] = [] + -- f (x ∷ xs) = x ∷ x ∷ xs + | Postulate {} + -- ^ e.g. `postulate pred : β„• β†’ β„•` + | Primitive {} + -- ^ e.g. `primitive primShowNat : β„• β†’ String` + deriving (Generic, Show, ToJSON, FromJSON) + +-- | Function clauses. +data Clause = Clause + { _telescope :: Telescope + -- ^ the telescope induced by the clause's context and patterns + , patterns :: [Pattern] + -- ^ the actual patterns of this function clause + , body :: Maybe Term + -- ^ the right hand side of the clause (@Nothing@ for absurd clauses) + } deriving (Generic, Show) +instance ToJSON Clause where toJSON = genericToJSON jsonOpts +instance FromJSON Clause where parseJSON = genericParseJSON jsonOpts + +-- | A telescope is a sequence of (named) types, a.k.a. bindings. +type Telescope = [Named (Pretty Type)] +-- | We under-approximate patterns as terms, +-- e.g. losing information about forced patterns. +type Pattern = Term +-- | Types are the same as terms. +type Type = Term + +-- | The AST of Agda terms. +data Term + = Pi Bool (Named Term) Term -- ^ e.g. @βˆ€ {A : Set}. A β†’ A@ + | Lam (Named Term) -- ^ e.g. @Ξ» x. x@ + | App Head [Term] -- ^ e.g. @f x (x + x)@ or @@0 (Ξ» x. x)@ + | Lit String -- ^ e.g. @42@ or @"something"@ + | Sort String -- ^ e.g. @Set@ + | Level String -- ^ e.g. @0β„“@ + | UnsolvedMeta -- ^ i.e. @{!!}@ + deriving (Generic, Show, FromJSON) + +instance ToJSON Term where + toJSON = \case + (Pi isDep (n :~ dom) codom) -> object + [ tag "Pi" + , "name" .= toJSON n -- T0D0: remove if "(nothing) (or _)?" + , "domain" .= toJSON dom + , "codomain" .= toJSON codom + ] + (Lam (n :~ f)) -> object + [ tag "Lambda" + , "abstraction" .= toJSON n + , "body" .= toJSON f + ] + (App f xs) -> + if null xs then + refHead + else + object [tag "Application", "head" .= refHead, "arguments" .= toJSON xs] + where + refHead = object $ case f of + (Left n) -> [tag "ScopeReference", "name" .= toJSON n] + (Right i) -> [tag "DeBruijn", "index" .= toJSON i] + (Lit s) -> object [tag "Literal", "literal" .= toJSON s] + (Sort s) -> object [tag "Sort", "sort" .= toJSON s] + (Level s) -> object [tag "Level", "level" .= toJSON s] + UnsolvedMeta -> object [tag "UnsolvedMetavariable"] + where tag s = "tag" .= JSON.String s + +-- * Conversion from Agda's internal syntax + +-- | Converting between two types @a@ and @b@ under Agda's typechecking monad. +-- +-- NB: 'go' is only used internally to de-clutter the recursive calls. +class (~>) a b | a -> b where + convert, go :: a -> TCM b + convert = go + +instance A.Definition ~> Definition where + go = go . theDef + +instance A.Defn ~> Definition where + go = \case + A.AbstractDefn defn -> go defn + A.Function{..} -> let cls = takeWhile isNotCubical funClauses in + -- NB: handle funWith and funExtLam + Function <$> traverse go cls + A.Datatype{..} -> do + -- NB: what is a dataClause??? + tys <- fmap unEl <$> traverse typeOfConst dataCons + ADT <$> traverse go tys + A.Record{..} -> do + -- NB: incorporate conHead/namedCon in the future for accuracy + -- + to solve the issue with private (non-public) fields + (tel, fs) <- splitAt recPars <$> go recTel + return $ Record tel (thing . item <$> fs) + A.Constructor{..} -> do + let cn = conName conSrcCon + d <- theDef <$> getConstInfo conData + return $ case d of + A.Datatype{..} -> + let Just ix = elemIndex (unqualify cn) (unqualify <$> dataCons) + in Constructor (pp conData) (toInteger ix) + A.Record{..} -> Constructor (pp conData) 0 + A.Primitive{..} -> return Primitive + A.PrimitiveSort{..} -> return Primitive + A.Axiom{..} -> return Postulate + d@A.DataOrRecSig{..} -> panic "dataOrRecSig" d + d@A.GeneralizableVar -> panic "generalizable variable" d + +instance A.Clause ~> Clause where + go A.Clause{..} = + Clause <$> go clauseTel + <*> traverse go (A.namedThing . unArg <$> namedClausePats) + -- ^ drop visibility and name information + <*> traverse go clauseBody + +instance A.DeBruijnPattern ~> Pattern where + go = \case + A.VarP _ v -> return $ App (Right $ dbPatVarIndex v) [] + A.DotP _ t -> go t + A.ConP c _ ps -> do + App (Left $ pp c) <$> traverse go (A.namedThing . unArg <$> ps) + A.LitP _ lit -> return $ Lit (pp lit) + A.ProjP _ qn -> return $ App (Left $ pp qn) [] + p@(A.IApplyP _ _ _ _) -> panic "pattern (cubical)" p + p@(A.DefP _ _ _) -> panic "pattern (cubical)" p + +instance A.Telescope ~> Telescope where + go = traverse action . A.telToList + where + action :: A.Dom (Name, A.Type) -> TCM (Named (Pretty Type)) + action dty = do + let (n, ty) = unDom dty + pty <- ppm ty + ty' <- go ty + let pdty = prender $ pDom dty $ P.text $ n <> " : " <> prender pty + return $ n :~ pdty :> ty' + +instance A.Type ~> Type where + go = go . A.unEl + +instance A.Term ~> Term where + go = flip (.) A.unSpine $ \case + -- ** abstractions + (A.Pi ty ab) -> do + let nameUsed = pp (A.domName ty) `notElem` ["_", "(nothing)"] + ty' <- go (unEl $ unDom ty) + ab' <- go (unEl $ unAbs ab) + return $ Pi nameUsed (absName ab :~ ty') ab' + (A.Lam _ ab) -> do + ab' <- go (unAbs ab) + return $ Lam (pp (absName ab) :~ ab') + -- ** applications + (A.Var i xs) -> App (Right i) <$> (traverse go xs) + (A.Def f xs) -> App (Left $ ppName f) <$> (traverse go xs) + (A.Con c _ xs) -> App (Left $ ppName $ conName c) <$> (traverse go xs) + -- ** other constants + (A.Lit x) -> return $ Lit $ pp x + (A.Level x) -> return $ Level $ pp x + (A.Sort x) -> return $ Sort $ pp x + (A.MetaV _ _) -> return UnsolvedMeta + -- ** there are some occurrences of `DontCare` in the standard library + (A.DontCare t) -> go t + -- ** crash on the rest (should never be encountered) + t@(A.Dummy _ _) -> panic "term" t + +instance A.Elim ~> Term where + go = \case + (A.Apply x) -> go (unArg x) + (A.Proj _ qn) -> return $ App (Left $ ppName qn) [] + (A.IApply _ _ x) -> go x + +-- * Utilities + +pp :: P.Pretty a => a -> String +pp = P.prettyShow + +ppm :: (MonadPretty m, P.PrettyTCM a) => a -> m Doc +ppm = P.prettyTCM + +prender :: Doc -> String +prender = P.renderStyle (P.Style P.OneLineMode 0 0.0) + +pinterleave :: (Applicative m, Semigroup (m Doc)) => m Doc -> [m Doc] -> m Doc +pinterleave sep = fsep . punctuate sep + +pbindings :: (MonadPretty m, PrettyTCM a) => [(String, a)] -> [m Doc] +pbindings = map $ \(n, ty) -> parens $ ppm n <> " : " <> ppm ty + +report :: MonadTCM m => VerboseLevel -> TCM Doc -> m () +report n x = liftTCM $ reportSDoc "agda2train" n x + +panic :: (P.Pretty a, Show a) => String -> a -> b +panic s t = error $ + "[PANIC] unexpected " <> s <> ": " <> pp t <> "\n show: " <> pp (show t) + +ppName :: A.QName -> String +ppName qn = pp qn <> "<" <> show (fromEnum $ nameId $ qnameName qn) <> ">" + +unqualify :: A.QName -> String +unqualify = pp . qnameName + +(\/) :: (a -> Bool) -> (a -> Bool) -> (a -> Bool) +(f \/ g) x = f x || g x + +isNotCubical :: A.Clause -> Bool +isNotCubical A.Clause{..} + | Just (A.Def qn _) <- clauseBody + = pp (qnameModule qn) /= "Agda.Primitive.Cubical" + | otherwise + = True + +-- | Configure JSON to omit empty (optional) fields and switch +-- from camelCase to kebab-case. +jsonOpts :: JSON.Options +jsonOpts = defaultOptions + { omitNothingFields = True + , fieldLabelModifier = \case + ('_' : s) -> s + "scopeGlobal" -> "scope-global" + "scopeLocal" -> "scope-local" + "scopePrivate" -> "scope-private" + s -> s + } + +instance P.PrettyTCM A.Definition where + prettyTCM d = go (theDef d) + where + go = \case + A.AbstractDefn defn -> go defn + A.Function{..} -> let cls = takeWhile isNotCubical funClauses in + fsep $ punctuate " |" + $ ppm . NamedClause (defName d) True <$> cls + A.Datatype{..} -> do + tys <- fmap unEl <$> traverse typeOfConst dataCons + pinterleave " |" $ pbindings $ zip (unqualify <$> dataCons) tys + A.Record{..} -> + let (tel, fs) = splitAt recPars $ telToList recTel in + (if null tel then "" else ppm (telFromList tel) <> " |- ") + <> (braces $ pinterleave " ;" $ pbindings $ unDom <$> fs) + A.Constructor{..} -> do + let cn = conName conSrcCon + d <- theDef <$> getConstInfo conData + case d of + A.Datatype{..} -> + let Just ix = elemIndex (unqualify cn) (unqualify <$> dataCons) + in ppm conData <> "@" <> ppm ix + A.Record{..} -> ppm conData <> "@0" + A.Primitive{..} -> "" + A.PrimitiveSort{..} -> "" + A.Axiom{..} -> "" + A.DataOrRecSig{..} -> "" + A.GeneralizableVar -> "" diff --git a/src/ToTrain.hs b/src/ToTrain.hs new file mode 100644 --- /dev/null +++ b/src/ToTrain.hs @@ -0,0 +1,186 @@ +{-# LANGUAGE TemplateHaskell #-} +-- | This module contains everything related to the generation of the training data. +module ToTrain where + +import Data.Maybe ( isJust ) +import Data.List ( isPrefixOf, isInfixOf, find, nub ) +import qualified Data.Set as S +import Data.FileEmbed ( embedStringFile ) + +import Control.Monad ( forM_, void, when, unless ) +import Control.Monad.Writer ( WriterT(runWriterT) ) +import Control.Monad.Error.Class ( catchError ) +import Control.Monad.IO.Class ( liftIO ) +import Control.Concurrent ( threadDelay ) +import Control.Concurrent.Async ( race ) + +import Agda.Syntax.Common + ( unArg, defaultArg, defaultArgInfo, namedArg ) +import Agda.Syntax.Internal +import Agda.Syntax.Internal.Names ( namesIn ) +import Agda.Syntax.Literal ( Literal(..) ) +import Agda.Syntax.Internal.Generic ( TermLike, foldTerm ) +import Agda.Syntax.Scope.Base ( nsInScope, allThingsInScope ) +import Agda.Syntax.Scope.Monad ( getCurrentScope ) + +import Agda.Utils.Monad ( whenM, tell1 ) +import Agda.Utils.Either ( caseEitherM ) + +import Agda.TypeChecking.Monad hiding (Reduced) +import Agda.TypeChecking.Reduce + ( Simplify, simplify, Normalise, normalise, Reduce, reduce ) +import Agda.TypeChecking.CheckInternal ( Action(..), defaultAction, checkInternal' ) +import Agda.TypeChecking.Pretty ( PrettyTCM ) + + +import AgdaInternals () +import Output hiding ( Definition(..), Clause(..), Term(..), Type(..) ) + +-- * Wrapper around Agda's typechecking monad 'TCM' + +-- | Additionally records/outputs training samples. +type C = WriterT [Sample] TCM + +runC :: C () -> TCM [Sample] +runC = (snd <$>) . runWriterT + +noop :: C () +noop = return () + +silently :: C a -> C () +silently k = void k `catchError` \ _ -> noop + +-- * Training data generation + +-- | A training function generates training data for each typed sub-term, +-- with access to the local context via the typechecking monad. +type TrainF = Type -> Term -> C () + +-- | An example training function that records a 'Output.Sample' +-- (i.e. context, type, and term) for a given subterm. +train :: TrainF +train ty t = do + let ns = names t + allNs <- nsInScope . allThingsInScope <$> liftTCM getCurrentScope + unless (null ns) $ + when (S.fromList ns `S.isSubsetOf` allNs) $ do + ctx <- getContextTelescope + pctx <- liftTCM $ ppm ctx; pty <- liftTCM $ ppm ty; pt <- liftTCM $ ppm t + rty <- mkReduced ty + report 30 $ "rty: " <> ppm (original rty) + + rt <- mkReduced t + + ctx' <- liftTCM $ convert ctx + rty' <- liftTCM $ traverse convert rty + rt' <- liftTCM $ traverse convert rt + tell1 $ Sample + { ctx = prender pctx :> ctx' + , goal = prender pty :> rty' + , term = prender pt :> rt' + , premises = map ppName ns + } + report 20 "{" + report 20 $ " ctx: " <> ppm (pp ctx) + report 30 $ " *pretty: " <> pure pctx + report 20 $ " goal: " <> ppm (pp ty) + report 30 $ " *pretty: " <> pure pty + reportReduced rty + report 20 $ " term: " <> ppm (pp t) + report 30 $ " *pretty: " <> pure pt + reportReduced rt + report 20 $ " premises: " <> ppm ns + report 20 "}" + +-- | Run the training function on each subterm of a definition. +forEachHole :: TrainF -> Definition -> C () +forEachHole trainF def@Defn{..} = unless (ignoreDef def) $ do + report 10 $ "------ definition: " <> ppm (pp defName) <> " -------" + sc <- getScope + case theDef of + (Function{..}) -> + forM_ funClauses $ \c@(Clause{..}) -> addContext clauseTel $ + case (clauseBody, unArg <$> clauseType) of + (Just t, Just ty) -> go ty t + _ -> noop + _ -> noop + -- TODO: get data from other places other than clause bodies + where + ignoreDef :: Definition -> Bool + ignoreDef Defn{..} + = False + || tooSlow (pp defName) + -- || defCopy + -- || defNoCompilation + -- || null (inverseScopeLookupName defName sc) + -- || isAnonymousModuleName (qnameModule defName) + -- || ("._." `isInfixOf` pp defName) + -- || (getOrigin defName /= UserWritten) + -- || ( ("with-" `isPrefixOf` pp (qnameName defName)) + + ignore :: Type -> C Bool + ignore ty = do + ctx <- fmap (snd . unDom) <$> getContext + return (ignoreType ty || any ignoreCtxType ctx) + + ignoreType, ignoreCtxType :: Type -> Bool + ignoreType = any cubicalRelated . map pp . names . unEl + ignoreCtxType = any cubicalRelated . map pp . names . unEl + + cubicalRelated, tooSlow :: String -> Bool + cubicalRelated = ("Agda.Primitive.Cubical.I" `isInfixOf`) + tooSlow + = ("Data.Rational.Properties" `isPrefixOf`) + \/ ("Prelude.Solvers" `isPrefixOf`) + \/ ("foundation.partitions" `isPrefixOf`) + \/ (`elem` defsToSkip) + + go :: Type -> Term -> C () + go ty t = whenM (not <$> ignore ty) + $ silently (checkInternal' act t CmpLeq ty) + + act :: Action C + act = defaultAction {preAction = pre} + + pre :: Type -> Term -> C Term + pre ty t = trainF ty t >> return t + +-- | Read a list of definitions to skip from @data/defsToSkip.txt@. +defsToSkip :: [String] +defsToSkip = lines $ $(embedStringFile "data/defsToSkip.txt") + +-- | Gathering names from terms. +names :: Term -> [QName] +names = nub . namesIn + +-- * Reduction + +-- | The hard limit on how much time can be spent on normalising a single term. +maxDuration = 2 -- seconds + +withTimeout :: TCM a -> TCM (Maybe a) +withTimeout k = getTC >>= \ s -> liftIO $ + caseEitherM + (race (threadDelay (maxDuration * 1000000)) + (fst <$> runSafeTCM k s)) + (\() -> pure Nothing) + (pure . Just) + +mkReduced :: + ( MonadFail m, MonadTCM m, PrettyTCM a + , Simplify a, Reduce a, Normalise a, Eq a + ) => a -> m (Reduced a) +mkReduced t = do + -- try different reductions (with timeout) + [_, simplified, reduced, normalised] <- compressList . (Just t :) + <$> mapM (liftTCM . withTimeout) [simplify t, reduce t, normalise t] + return $ Reduced {original = t, ..} + where + compressList :: Eq a => [Maybe a] -> [Maybe a] + compressList xs = find ((`notElem` xs) . Just) <$> xs + +reportReduced :: (MonadTCM m, PrettyTCM a) => Reduced a -> m () +reportReduced Reduced{..} = do + report 30 $ " *simplified: " <> ppm simplified + report 30 $ " *reduced: " <> ppm reduced + report 30 $ " *normalised: " <> ppm normalised