diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Eric Conlon (c) 2021 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * 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. + + * Neither the name of Eric Conlon nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +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,15 @@ +# overeasy + +[![CircleCI](https://circleci.com/gh/ejconlon/overeasy/tree/master.svg?style=svg)](https://circleci.com/gh/ejconlon/overeasy/tree/master) + +A purely functional E-Graph library + +## How to build and run + +This repo is setup so you never have to `cd` out of the root directory. + +To run the Haskell programs, you need `stack` installed on your system. [This](https://docs.haskellstack.org/en/stable/README/) is an easy way to do so, but you can also check your package manager or use [ghcup](https://www.haskell.org/ghcup/). + +Most of the interesting stuff is going to be run in the test suite. Run it with `make test`. `stack` will get the appropriate Haskell compiler and package dependencies, and it will compile the project before running the test suite. + +If you have Docker installed and running (with about 4G of RAM allocated to it), you can run `make docker-test` to build and run tests in a container instead of installing `stack` locally. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/overeasy.cabal b/overeasy.cabal new file mode 100644 --- /dev/null +++ b/overeasy.cabal @@ -0,0 +1,136 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.34.4. +-- +-- see: https://github.com/sol/hpack + +name: overeasy +version: 0.1.0 +synopsis: A purely functional E-Graph library +description: Please see the README on GitHub at +category: Data Structures +homepage: https://github.com/ejconlon/overeasy#readme +bug-reports: https://github.com/ejconlon/overeasy/issues +author: Eric Conlon +maintainer: ejconlon@gmail.com +copyright: (c) 2021 Eric Conlon +license: BSD3 +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + +source-repository head + type: git + location: https://github.com/ejconlon/overeasy + +library + exposed-modules: + Overeasy.Assoc + Overeasy.EGraph + Overeasy.EquivFind + Overeasy.Matching + Overeasy.Source + Overeasy.Streams + Overeasy.Util + other-modules: + Paths_overeasy + hs-source-dirs: + src + default-extensions: + BangPatterns + ConstraintKinds + DeriveFunctor + DeriveFoldable + DeriveGeneric + DeriveTraversable + DerivingStrategies + DerivingVia + FlexibleContexts + FlexibleInstances + FunctionalDependencies + GADTs + GeneralizedNewtypeDeriving + LambdaCase + KindSignatures + MultiParamTypeClasses + PatternSynonyms + Rank2Types + ScopedTypeVariables + StandaloneDeriving + TemplateHaskell + TupleSections + TypeApplications + TypeOperators + TypeFamilies + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints -fno-warn-unused-top-binds -fwrite-ide-info -hiedir=.hie + build-depends: + algebraic-graphs >=0.5 && <0.7 + , base >=4.12 && <5 + , containers ==0.6.* + , deepseq ==1.4.* + , hashable >=1.3 && <1.5 + , int-like >=0.1.1 && <0.2 + , logict >=0.7 && <0.9 + , mtl ==2.2.* + , recursion-schemes ==5.2.* + , text >=1.2 && <2.1 + , transformers >=0.5 && <0.7 + , unfree ==0.1.* + , unordered-containers ==0.2.* + default-language: Haskell2010 + +test-suite overeasy-test + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: + Test.Overeasy.Arith + Test.Overeasy.BinTree + Paths_overeasy + hs-source-dirs: + test + default-extensions: + BangPatterns + ConstraintKinds + DeriveFunctor + DeriveFoldable + DeriveGeneric + DeriveTraversable + DerivingStrategies + DerivingVia + FlexibleContexts + FlexibleInstances + FunctionalDependencies + GADTs + GeneralizedNewtypeDeriving + LambdaCase + KindSignatures + MultiParamTypeClasses + PatternSynonyms + Rank2Types + ScopedTypeVariables + StandaloneDeriving + TemplateHaskell + TupleSections + TypeApplications + TypeOperators + TypeFamilies + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints -fno-warn-unused-top-binds -fwrite-ide-info -hiedir=.hie -threaded -rtsopts -with-rtsopts=-N + build-depends: + algebraic-graphs >=0.5 && <0.7 + , base >=4.12 && <5 + , containers ==0.6.* + , deepseq ==1.4.* + , hashable >=1.3 && <1.5 + , hedgehog ==1.0.* + , int-like >=0.1.1 && <0.2 + , logict >=0.7 && <0.9 + , mtl ==2.2.* + , overeasy + , prop-unit ==0.1.* + , recursion-schemes ==5.2.* + , text >=1.2 && <2.1 + , transformers >=0.5 && <0.7 + , unfree ==0.1.* + , unordered-containers ==0.2.* + default-language: Haskell2010 diff --git a/src/Overeasy/Assoc.hs b/src/Overeasy/Assoc.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/Assoc.hs @@ -0,0 +1,254 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | See 'Assoc'. +module Overeasy.Assoc + ( Assoc + , assocFwd + , assocBwd + , assocEquiv + , assocSize + , assocNew + , assocSingleton + , AssocInsertRes (..) + , assocInsertInc + , assocInsert + , assocFromList + , assocToList + , assocMember + , assocLookupByKey + , assocPartialLookupByKey + , assocLookupByValue + , assocPartialLookupByValue + , assocLookupRoot + , assocRoots + , assocLeaves + , assocMembers + , assocCanCompact + , assocCompactInc + , assocCompact + , assocRemoveAllInc + , assocRemoveAll + , assocUnion + , assocFootprint + ) where + +import Control.DeepSeq (NFData) +import Control.Monad.State.Strict (MonadState (..), State, modify') +import Data.Coerce (Coercible) +import Data.Foldable (foldl') +import Data.Hashable (Hashable) +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap +import Data.Maybe (fromJust) +import GHC.Generics (Generic) +import IntLike.Map (IntLikeMap) +import qualified IntLike.Map as ILM +import IntLike.Set (IntLikeSet) +import qualified IntLike.Set as ILS +import Overeasy.EquivFind (EquivAddRes (..), EquivFind, efAddInc, efBwd, efCanCompact, efCompactInc, efEquivs, efLeaves, + efLookupRoot, efMember, efMembers, efNew, efRemoveAllInc, efRoots, efSingleton, + efUnsafeAddLeafInc, efUnsafeMerge) + +-- | Associates keys and values in such a way that inserting +-- duplicate values induces equivalences on their keys. +-- Invariant: fwd and bwd maps contain only root keys. +data Assoc x a = Assoc + { assocFwd :: !(IntLikeMap x a) + -- ^ Map from id to element + , assocBwd :: !(HashMap a x) + -- ^ Map from element to id + , assocEquiv :: !(EquivFind x) + -- ^ Equivalence classes of ids + } deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | How many values are in the map? +assocSize :: Assoc x a -> Int +assocSize = ILM.size . assocFwd + +-- | Creates an empty assoc +assocNew :: Assoc x a +assocNew = Assoc ILM.empty HashMap.empty efNew + +-- | Creates a singleton assoc +assocSingleton :: (Coercible x Int, Hashable a) => x -> a -> Assoc x a +assocSingleton x a = Assoc (ILM.singleton x a) (HashMap.singleton a x) (efSingleton x) + +-- | The result of inserting into the assoc, if you're interested. +data AssocInsertRes x = + AssocInsertResUnchanged + | AssocInsertResCreated + | AssocInsertResUpdated + | AssocInsertResMerged !(IntLikeSet x) + deriving stock (Eq, Show) + +-- | Insert into the assoc (raw version) +assocInsertInc :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a) +assocInsertInc x a1 assoc@(Assoc fwd bwd equiv) = finalRes where + finalRes = + let (res, equiv') = efAddInc x equiv + in case res of + EquivAddResNewRoot -> insertRoot x equiv' + EquivAddResAlreadyLeafOf z -> updateRoot z + EquivAddResAlreadyRoot -> updateRoot x + updateRoot w = + -- w is existing root and is guaranteed to map to something + let a0 = ILM.partialLookup w fwd + in if a0 == a1 + -- the value has not changed, don't need to change assoc + then ((w, AssocInsertResUnchanged), assoc) + else + -- value has changed, need to check if it's fresh + case HashMap.lookup a1 bwd of + -- never seen; insert and return + Nothing -> + let fwd' = ILM.insert w a1 fwd + bwd' = HashMap.insert a1 w (HashMap.delete a0 bwd) + in ((w, AssocInsertResUpdated), Assoc fwd' bwd' equiv) + -- mapped to another set of nodes, merge + Just v -> + let (toKeep, toDelete, equiv') = efUnsafeMerge w v equiv + res = AssocInsertResMerged toDelete + in if toKeep == w + -- w wins + then + let fwd' = ILM.insert w a1 (ILM.delete v fwd) + bwd' = HashMap.insert a1 w (HashMap.delete a0 bwd) + in ((w, res), Assoc fwd' bwd' equiv') + -- v wins + else + let fwd' = ILM.delete w fwd + bwd' = HashMap.delete a0 bwd + in ((v, res), Assoc fwd' bwd' equiv') + insertRoot w equiv' = + -- w is new root that doesn't exist + case HashMap.lookup a1 bwd of + -- never seen; insert and return + Nothing -> + let fwd' = ILM.insert w a1 fwd + bwd' = HashMap.insert a1 w bwd + in ((w, AssocInsertResCreated), Assoc fwd' bwd' equiv') + Just v -> + let (toKeep, toDelete, equiv'') = efUnsafeMerge w v equiv' + res = AssocInsertResMerged toDelete + in if toKeep == w + -- w wins + then + let fwd' = ILM.insert w a1 (ILM.delete v fwd) + bwd' = HashMap.insert a1 w bwd + in ((w, res), Assoc fwd' bwd' equiv'') + -- v wins + else + let fwd' = ILM.delete w fwd + in ((v, res), Assoc fwd' bwd equiv'') + +-- | Insert into the assoc (the state version) +assocInsert :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> State (Assoc x a) (x, AssocInsertRes x) +assocInsert x a = state (assocInsertInc x a) + +-- | Build an assoc from a list of pairs +assocFromList :: (Coercible x Int, Ord x, Eq a, Hashable a) => [(x, a)] -> Assoc x a +assocFromList = foldl' (\assoc (x, a) -> snd (assocInsertInc x a assoc)) assocNew + +-- | Turn an assoc into a list of pairs (NOTE - emits only root keys!) +assocToList :: Coercible x Int => Assoc x a -> [(x, a)] +assocToList = ILM.toList . assocFwd + +-- | Is the given key in the assoc? +assocMember :: Coercible x Int => x -> Assoc x a -> Bool +assocMember x (Assoc _ _ equiv) = efMember x equiv + +-- | Forward lookup +assocLookupByKey :: Coercible x Int => x -> Assoc x a -> Maybe a +assocLookupByKey x (Assoc fwd _ equiv) = ILM.lookup (efLookupRoot x equiv) fwd + +-- | PARTIAL forward lookup +assocPartialLookupByKey :: Coercible x Int => x -> Assoc x a -> a +assocPartialLookupByKey x = fromJust . assocLookupByKey x + +-- | Backward lookup +assocLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> Maybe x +assocLookupByValue a = HashMap.lookup a . assocBwd + +-- | PARTIAL backward lookup +assocPartialLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> x +assocPartialLookupByValue a = flip (HashMap.!) a . assocBwd + +-- | Finds the root for the given key (id if not found) +assocLookupRoot :: Coercible x Int => x -> Assoc x a -> x +assocLookupRoot x = efLookupRoot x . assocEquiv + +-- | List all root (live, non-compactible) keys +assocRoots :: Coercible x Int => Assoc x a -> [x] +assocRoots = efRoots . assocEquiv + +-- | List all leaf (dead, compactible) keys +assocLeaves :: Coercible x Int => Assoc x a -> [x] +assocLeaves = efLeaves . assocEquiv + +-- | List all entries (root and leaf) +assocMembers :: Coercible x Int => Assoc x a -> [x] +assocMembers = efMembers . assocEquiv + +-- | Are there dead keys in the equiv from 'assocInsert'? +assocCanCompact :: Assoc x a -> Bool +assocCanCompact = efCanCompact . assocEquiv + +-- | Removes all dead keys in the equiv (raw version). +assocCompactInc :: Coercible x Int => Assoc x a -> (IntLikeMap x x, Assoc x a) +assocCompactInc assoc@(Assoc fwd bwd equiv) = + let replacements = efBwd equiv + assoc' = + if ILM.null replacements + then assoc + else let (_, equiv') = efCompactInc equiv in Assoc fwd bwd equiv' + in (replacements, assoc') + +-- | Removes all dead keys in the equiv (state version). +-- Returns map of dead leaf node -> live root node +assocCompact :: Coercible x Int => State (Assoc x a) (IntLikeMap x x) +assocCompact = state assocCompactInc + +-- | Removes the given keys from the assoc (raw version) +assocRemoveAllInc :: (Coercible x Int, Eq a, Hashable a) => [x] -> Assoc x a -> Assoc x a +assocRemoveAllInc xs (Assoc fwd0 bwd0 equiv0) = Assoc fwdFinal bwdFinal equivFinal where + (remap, equivFinal) = efRemoveAllInc xs equiv0 + (fwdFinal, bwdFinal) = foldl' go (fwd0, bwd0) xs + go tup@(fwd, bwd) x = + case ILM.lookup x fwd of + -- Leaf, ignore + Nothing -> tup + -- Root + Just a -> + case ILM.lookup x remap of + -- Singleton root, delete + Nothing -> + let fwd' = ILM.delete x fwd + bwd' = HashMap.delete a bwd + in (fwd', bwd') + -- Remapped root, rotate + Just y -> + let fwd' = ILM.delete x (ILM.insert y a fwd) + bwd' = HashMap.insert a y bwd + in (fwd', bwd') + +-- | Removes the given keys from the assoc (state version). +-- Values will only be removed from the assoc if the key is a singleton root. +-- If a key is not found, it is simply ignored. +assocRemoveAll :: (Coercible x Int, Eq a, Hashable a) => [x] -> State (Assoc x a) () +assocRemoveAll = modify' . assocRemoveAllInc + +-- | Join two assocs (uses the first as the base) +assocUnion :: (Coercible x Int, Ord x, Eq a, Hashable a) => Assoc x a -> Assoc x a -> Assoc x a +assocUnion base (Assoc fwd _ equiv) = Assoc fwdFinal bwdFinal equivFinal where + goRoots assocGo (x, a) = snd (assocInsertInc x a assocGo) + goLeaves equivGo (leaf, oldRoot) = efUnsafeAddLeafInc oldRoot leaf equivGo + Assoc fwdFinal bwdFinal equivMid = foldl' goRoots base (ILM.toList fwd) + equivFinal = foldl' goLeaves equivMid (ILM.toList (efBwd equiv)) + +-- | Returns the footprint of the given value - all keys that map to it (root and leaf) +assocFootprint :: (Coercible x Int, Eq a, Hashable a) => a -> Assoc x a -> IntLikeSet x +assocFootprint a (Assoc _ bwd equiv) = + case HashMap.lookup a bwd of + Nothing -> ILS.empty + Just r -> efEquivs r equiv diff --git a/src/Overeasy/EGraph.hs b/src/Overeasy/EGraph.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/EGraph.hs @@ -0,0 +1,568 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | See 'EGraph'. +module Overeasy.EGraph + ( EClassId (..) + , ENodeId (..) + , EAnalysis + , noAnalysis + , EClassInfo (..) + , EGraph + , WorkItem + , WorkList + , ClassReplacements + , MergeResult (..) + , egClassSource + , egNodeSource + , egEquivFind + , egClassMap + , egNodeAssoc + , egHashCons + , egClassSize + , egNodeSize + , egFindNode + , egFindTerm + , egClassInfo + , egNew + , egClasses + , egCanonicalize + , egCanonicalizePartial + , egAddTerm + , egMerge + , egMergeMany + , egReanalyzeSubset + , egReanalyze + ) where + +import Control.DeepSeq (NFData) +import Control.Monad (unless, void) +import Control.Monad.State.Strict (State, gets, modify', state) +import Data.Foldable (foldl', toList) +import Data.Functor.Foldable (project) +import Data.Hashable (Hashable) +import Data.List.NonEmpty (NonEmpty (..)) +import Data.Maybe (maybeToList) +import Data.Semigroup (sconcat) +import Data.Sequence (Seq (..)) +import qualified Data.Sequence as Seq +import GHC.Generics (Generic) +import IntLike.Map (IntLikeMap (..)) +import qualified IntLike.Map as ILM +import IntLike.Set (IntLikeSet (..)) +import qualified IntLike.Set as ILS +import Overeasy.Assoc (Assoc, AssocInsertRes (..), assocCompactInc, assocFwd, assocInsertInc, assocLookupByValue, + assocMember, assocMembers, assocNew, assocPartialLookupByKey, assocRemoveAllInc, assocSingleton, + assocUnion) +import Overeasy.EquivFind (EquivFind (..), EquivMergeSetsRes (..), efAddInc, efCanonicalize, efCanonicalizePartial, + efClosure, efCompactInc, efFindRoot, efLookupRoot, efMergeSetsInc, efNew, efRootsSize, + efSubset) +import Overeasy.Source (Source, sourceAddInc, sourceNew) +import Overeasy.Util (Changed (..), RecursiveWhole, foldWholeM, stateFold) + +-- | An opaque class id. +-- Constructor exported for coercibility. +-- Num instance for literals only. +newtype EClassId = EClassId { unEClassId :: Int } + deriving stock (Show) + deriving newtype (Eq, Ord, Enum, Hashable, NFData, Num) + +-- | An opaque node id +-- Constructor exported for coercibility. +-- Num instance for literals only. +newtype ENodeId = ENodeId { unENodeId :: Int } + deriving stock (Show) + deriving newtype (Eq, Ord, Enum, Hashable, NFData, Num) + +-- | The definition of an 'EGraph' analysis. +-- 'd' must be a join semilattice. +-- This function must be monotonic. +type EAnalysis d f = f d -> d + +-- | A disabled analysis +noAnalysis :: EAnalysis () f +noAnalysis = const () + +-- private +-- An internal triple of node, class, and data +data ENodeTriple d = ENodeTriple + { entNode :: !ENodeId + , entClass :: !EClassId + , entData :: !d + } deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Info stored for every class: analysis data, class members (nodes), and parent nodes. +data EClassInfo d f = EClassInfo + { eciData :: !d + , eciNodes :: !(Assoc ENodeId (f ())) + , eciParents :: !(IntLikeSet ENodeId) + } deriving stock (Generic) + +deriving stock instance (Eq d, Eq (f ())) => Eq (EClassInfo d f) +deriving stock instance (Show d, Show (f ())) => Show (EClassInfo d f) +deriving anyclass instance (NFData d, NFData (f ())) => NFData (EClassInfo d f) + +-- | A set of class ids to merge +type WorkItem = IntLikeSet EClassId + +-- | A sequences of groups of class ids to mrege +type WorkList = Seq WorkItem + +-- | An invertible multimap of new root class to the sets of old classes it subsumes +-- Can be used to externally recanonicalize any structures that reference class ids +-- after merges. +type ClassReplacements = EquivFind EClassId + +-- | Merging classes can result in a few outcomes: +data MergeResult a = + MergeResultUnchanged + -- ^ All classes already merged, no change + | MergeResultMissing !WorkItem + -- ^ Some classes missing, returns first problematic worklist entry + -- (note that not all classes in worklist item will be missing, + -- only at least one from the set) + | MergeResultChanged !a + -- ^ Some classes merged, returns root map or merged class id + deriving stock (Eq, Show, Functor, Foldable, Traversable) + +-- | An E-Graph implementation +data EGraph d f = EGraph + { egClassSource :: !(Source EClassId) + -- ^ Id source for classes + , egNodeSource :: !(Source ENodeId) + -- ^ Id source for nodes + , egEquivFind :: !(EquivFind EClassId) + -- ^ Class equivalences + , egClassMap :: !(IntLikeMap EClassId (EClassInfo d f)) + -- ^ Map of class to info + -- Invariant: Only contains root classes. + , egNodeAssoc :: !(Assoc ENodeId (f EClassId)) + -- ^ Assoc of node id to node structure + -- Invariant: only contains canonical structures (with root classes). + , egHashCons :: !(IntLikeMap ENodeId EClassId) + -- ^ Map of node to class + -- Invariant: only contains root classes. + } deriving stock (Generic) + +deriving stock instance (Eq d, Eq (f EClassId), Eq (f ())) => Eq (EGraph d f) +deriving stock instance (Show d, Show (f EClassId), Show (f ())) => Show (EGraph d f) +deriving anyclass instance (NFData d, NFData (f EClassId), NFData (f ())) => NFData (EGraph d f) + +-- | Number of equivalent classes in the 'EGraph' (see 'ufSize') +egClassSize :: EGraph d f -> Int +egClassSize = efRootsSize . egEquivFind + +-- | Number of nodes in the 'EGraph' +egNodeSize :: EGraph d f -> Int +egNodeSize = ILM.size . egHashCons + +-- | Lookup info for the given 'EClass' +egClassInfo :: EClassId -> EGraph d f -> Maybe (EClassInfo d f) +egClassInfo c = ILM.lookup c . egClassMap + +-- | Find the class of the given node, if it exists. +-- Note that you may have to canonicalize first to find it! +egFindNode :: (Eq (f EClassId), Hashable (f EClassId)) => f EClassId -> EGraph d f -> Maybe EClassId +egFindNode fc eg = do + n <- assocLookupByValue fc (egNodeAssoc eg) + ILM.lookup n (egHashCons eg) + +-- | Find the class of the given term, if it exists +egFindTerm :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId)) => t -> EGraph d f -> Maybe EClassId +egFindTerm t eg = foldWholeM (`egFindNode` eg) t + +-- | Creates a new 'EGraph' +egNew :: EGraph d f +egNew = EGraph (sourceNew (EClassId 0)) (sourceNew (ENodeId 0)) efNew ILM.empty assocNew ILM.empty + +-- | Yields all root classes +egClasses :: State (EGraph d f) (IntLikeSet EClassId) +egClasses = gets (ILM.keysSet . egClassMap) + +-- | Find the canonical form of a node. +-- If any classes are missing, the first missing is returned. +egCanonicalize :: Traversable f => f EClassId -> State (EGraph d f) (Either EClassId (f EClassId)) +egCanonicalize fc = gets (efCanonicalize fc . egEquivFind) + +-- | Find the canonical form of a node. +-- If any classes are missing, simply skip them. +egCanonicalizePartial :: Traversable f => f EClassId -> State (EGraph d f) (f EClassId) +egCanonicalizePartial fc = gets (efCanonicalizePartial fc . egEquivFind) + +-- private +-- Variant of canonicalization used in rebuilding - updates node assocs and returns +-- 1. New canonical node id (could be the same) +-- 2. Maybe a set of node ids to delete (no longer canonical) +egCanonicalizeInternal :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => ENodeId -> State (EGraph d f) (ENodeId, Maybe (IntLikeSet ENodeId)) +egCanonicalizeInternal x = state $ \eg -> + let ef = egEquivFind eg + assoc = egNodeAssoc eg + node = assocPartialLookupByKey x assoc + fz = efCanonicalizePartial node ef + ((y, res), assoc') = assocInsertInc x fz assoc + in case res of + AssocInsertResUnchanged -> ((y, Nothing), eg) + AssocInsertResMerged toDelete -> + ((y, Just toDelete), eg { egNodeAssoc = assoc' }) + _ -> ((y, Nothing), eg { egNodeAssoc = assoc' }) + +-- private +data AddNodeRes d = AddNodeRes !Changed !(Seq (ENodeTriple d)) + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +instance Semigroup (AddNodeRes d) where + AddNodeRes c1 p1 <> AddNodeRes c2 p2 = AddNodeRes (c1 <> c2) (p1 <> p2) + +instance Monoid (AddNodeRes d) where + mempty = AddNodeRes ChangedNo Seq.empty + mappend = (<>) + +-- private +egAddNodeSub :: (Functor f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> f (ENodeTriple d) -> State (EGraph d f) (Changed, ENodeTriple d) +egAddNodeSub ana fcd = do + let fc = fmap entClass fcd + -- important: node should already be canonicalized! + -- first lookup the node in the assoc to ensure uniqueness + mayNodeId <- gets (assocLookupByValue fc . egNodeAssoc) + case mayNodeId of + Just n -> do + x <- gets (ILM.partialLookup n . egHashCons) + eci <- gets (ILM.partialLookup x . egClassMap) + let d = eciData eci + pure (ChangedNo, ENodeTriple n x d) + Nothing -> state $ \eg -> + -- node does not exist; get new node and class ids + let (n, nodeSource') = sourceAddInc (egNodeSource eg) + (x, classSource') = sourceAddInc (egClassSource eg) + -- add it to the uf (can discard return value since it's a new id, will be the same) + (_, ef') = efAddInc x (egEquivFind eg) + -- add it to the assoc (ignore and partial by construction) + (_, assoc') = assocInsertInc n fc (egNodeAssoc eg) + -- insert into the hashcons + hc' = ILM.insert n x (egHashCons eg) + -- analyze the node and put that info in the class map + d = ana (fmap entData fcd) + eci = EClassInfo d (assocSingleton n (void fcd)) ILS.empty + classMap' = ILM.insert x eci (egClassMap eg) + eg' = eg { egNodeSource = nodeSource', egClassSource = classSource', egEquivFind = ef', egNodeAssoc = assoc', egHashCons = hc', egClassMap = classMap' } + in ((ChangedYes, ENodeTriple n x d), eg') + +-- private +egAddTermSub :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> t -> State (EGraph d f) (AddNodeRes d, ENodeTriple d) +egAddTermSub ana = go where + go t = do + -- unwrap to work with the functor layer + let ft = project t + -- add all child nodes + frx <- traverse go ft + -- collect info generated from child nodes and leave pure structure + let (AddNodeRes changed1 children, fx) = sequenceA frx + -- now fx should be canonicalized by construction + -- add the node to get its node and class ids + (changed2, z@(ENodeTriple n _ _)) <- egAddNodeSub ana fx + -- now update all its children to add this as a parent + unless (Seq.null children) $ + modify' $ \eg -> + -- Add node to class parents (unless it's a self parent) + let cm' = foldl' (\cm (ENodeTriple _ c _) -> ILM.adjust (\v -> if assocMember n (eciNodes v) then v else v { eciParents = ILS.insert n (eciParents v) }) c cm) (egClassMap eg) children + in eg { egClassMap = cm' } + pure (AddNodeRes (changed1 <> changed2) (Seq.singleton z), z) + +-- | Adds a term (recursively) to the graph. If already in the graph, returns 'ChangedNo' and existing class id. Otherwise +-- returns 'ChangedYes' and a new class id. +egAddTerm :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> t -> State (EGraph d f) (Changed, EClassId) +egAddTerm ana t = fmap (\(AddNodeRes c _, ENodeTriple _ x _) -> (c, x)) (egAddTermSub ana t) + +-- | Merges two classes: +-- Returns 'Nothing' if the classes are not found or if they're already equal. +-- Otherwise returns the class remapping. +-- Note that it's MUCH more efficient to accumulate a 'WorkList' and use 'egMergeMany'. +egMerge :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ())) + => EClassId -> EClassId -> State (EGraph d f) (MergeResult EClassId) +egMerge i j = do + mr <- egMergeMany (Seq.singleton (ILS.fromList [i, j])) + -- We're guaranteed to have one and only one root in the map, so this won't fail + pure (fmap (fst . head . ILM.toList . efFwd . fst) mr) + +-- private +data BuildWorkResult a = + BuildWorkResultUnchanged + | BuildWorkResultMissing !WorkItem + | BuildWorkResultOk !a + +-- private +egBuildWorkItem :: WorkItem -> State (EGraph d f) (BuildWorkResult WorkItem) +egBuildWorkItem cs = do + mayRoots <- fmap (\ef -> traverse (`efFindRoot` ef) (ILS.toList cs)) (gets egEquivFind) + pure $! case mayRoots of + Nothing -> BuildWorkResultMissing cs + Just roots -> + let rootsSet = ILS.fromList roots + in if ILS.size rootsSet < 2 + then BuildWorkResultUnchanged + else BuildWorkResultOk rootsSet + +-- private +egBuildWorklist :: WorkList -> State (EGraph d f) (BuildWorkResult WorkList) +egBuildWorklist = go Empty where + go !acc = \case + Empty -> + pure $! if Seq.null acc + then BuildWorkResultUnchanged + else BuildWorkResultOk acc + cs :<| wl' -> do + rcs <- egBuildWorkItem cs + case rcs of + BuildWorkResultUnchanged -> go acc wl' + BuildWorkResultMissing cs' -> pure (BuildWorkResultMissing cs') + BuildWorkResultOk cs' -> go (acc :|> cs') wl' + +-- | Merges many sets of classes. +-- Returns 'Nothing' if the classes are not found or if they're already equal. +-- Otherwise returns the class remapping (equiv map of root to set of leaf classes). +-- It is important to note that the leaf classes in the returned mapping have been +-- REMOVED from the egraph, so they cannot be used to lookup classes in the future. +-- Therefore, if you have any class ids stored externally, you'll want to (partially) +-- canonicalize with the returned mapping. +-- Also note that the analysis of a given class is going to be an UNDER-APPROXIMATION +-- of the true analysis value, because per-node analyses are not recomputed. +egMergeMany :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ())) + => WorkList -> State (EGraph d f) (MergeResult (ClassReplacements, IntLikeSet EClassId)) +egMergeMany wl0 = do + br <- egBuildWorklist wl0 + case br of + BuildWorkResultUnchanged -> pure MergeResultUnchanged + BuildWorkResultMissing cs -> pure (MergeResultMissing cs) + BuildWorkResultOk wl1 -> fmap MergeResultChanged (egRebuild wl1) + +-- private +-- Folds over items in worklist to merge, returning: +-- 1. map of old class -> new class for changed classes only +-- 2. closure of remapped classes (includes roots) +egRebuildMerge :: WorkList -> State (EGraph d f) (IntLikeMap EClassId EClassId, IntLikeSet EClassId) +egRebuildMerge wl = finalRes where + finalRes = state $ \eg -> + let ef = egEquivFind eg + in case efMergeSetsInc (toList wl) ef of + EquivMergeSetsResChanged roots classRemapSet ef' -> + let classRemap = ILM.fromList (fmap (\c -> (c, efLookupRoot c ef')) (ILS.toList classRemapSet)) + closure = efClosure (ILS.toList roots) ef' + in ((classRemap, closure), eg { egEquivFind = ef' }) + _ -> ((ILM.empty, ILS.empty), eg) + +-- private +-- Loop through nodes of all changed classes and update the hashcons to point to new classes +egRebuildHashCons :: IntLikeMap EClassId EClassId -> State (EGraph d f) () +egRebuildHashCons classRemap = + modify' (\eg -> let hc' = foldl' (go (egClassMap eg)) (egHashCons eg) (ILM.toList classRemap) in eg { egHashCons = hc' }) where + go cm hc (oldClassId, newClassId) = + let eci = ILM.partialLookup oldClassId cm + nodes = eciNodes eci + in foldl' (flip (`ILM.insert` newClassId)) hc (assocMembers nodes) + +-- private +-- For each touched class, recanonicalize all its nodes +-- Return pair of +-- 1. Set of parent class ids that can observe changes (i.e. need recanonicalization/reanalysis) +-- 2. Worklist of induced parent equalities found by recanonicalization +egRebuildAssoc :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => IntLikeMap ENodeId EClassId -> IntLikeMap EClassId EClassId -> IntLikeSet EClassId -> State (EGraph d f) (IntLikeSet EClassId, WorkList) +egRebuildAssoc origHc classRemap touchedClasses = do + hc <- gets egHashCons + cm <- gets egClassMap + -- For each class that we're going to merge + stateFold (ILS.empty, Empty) (ILS.toList touchedClasses) $ \(ps, parentWl) c -> do + -- Get the class info + let eci = ILM.partialLookup c cm + -- For each node in the class + (finalChanged, finalParentWl) <- stateFold (False, parentWl) (assocMembers (eciNodes eci)) $ \(changed', parentWl') n -> do + -- Canonicalize it and add to the node map + (newN, mayEquivNs) <- egCanonicalizeInternal n + case mayEquivNs of + Nothing -> pure (changed', parentWl') + Just equivNs -> + let allNs = ILS.insert newN equivNs + allEquivClasses = ILS.map (`ILM.partialLookup` hc) allNs + in if ILS.size allEquivClasses > 1 + then pure (True, parentWl' :|> allEquivClasses) + else pure (changed', parentWl') + -- Emit observing parents if: + -- 1. class has changed + -- 2. any nodes have changed during canonicalization + -- Note that we look up parents in the ORIGINAL hashcons because those are the ones that have the nodes pointing to this + let emitParents = finalChanged || ILM.member c classRemap + addlParents = ILS.map (`ILM.partialLookup` origHc) (eciParents eci) + ps' = if emitParents then ILS.union addlParents ps else ps + pure (ps', finalParentWl) + +-- private +-- One round of rebuilding +egRebuildNodeRound :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => IntLikeMap ENodeId EClassId -> WorkList -> IntLikeSet EClassId -> State (EGraph d f) (IntLikeSet EClassId, WorkList, IntLikeSet EClassId) +egRebuildNodeRound origHc wl parents = do + -- First merge all classes together and get merged class sets + (classRemap, classClosure) <- egRebuildMerge wl + -- Now update the hashcons so node ids point to merged classes + egRebuildHashCons classRemap + -- Track all classes touched here + let touchedClasses = ILS.union parents classClosure + -- Traverse all touched classes and canonicalize their nodes, + -- recording the mapping from old -> new + -- Also track parents that can observe changes to this class + (candParents, parentWl) <- egRebuildAssoc origHc classRemap touchedClasses + -- (We ignore parents that we have just now rebuilt) + let finalParents = ILS.difference candParents touchedClasses + pure (touchedClasses, parentWl, finalParents) + +-- private +-- Rebuild just the class info corresponding to 'newClass' +egRebuildClassSingle :: (Semigroup d, Eq (f ()), Hashable (f ())) => EClassId -> IntLikeSet EClassId -> IntLikeMap EClassId (EClassInfo d f) -> IntLikeMap EClassId (EClassInfo d f) +egRebuildClassSingle newClass oldClasses initCm = + let EClassInfo rootData rootNodes rootParents = ILM.partialLookup newClass initCm + finalData = sconcat (rootData :| fmap (\c -> eciData (ILM.partialLookup c initCm)) (ILS.toList oldClasses)) + -- keep dead self nodes here. will be dropped in compact + finalNodes = foldl' (\s c -> assocUnion s (eciNodes (ILM.partialLookup c initCm))) rootNodes (ILS.toList oldClasses) + -- keep dead parent nodes here, just exclude self nodes. will be dropped in compact + lookupParents c = eciParents (ILM.partialLookup c initCm) + candParents = foldl' (\s c -> ILS.union s (lookupParents c)) rootParents (ILS.toList oldClasses) + finalParents = ILS.difference candParents (ILS.fromList (assocMembers finalNodes)) + finalInfo = EClassInfo finalData finalNodes finalParents + finalCm = ILM.insert newClass finalInfo initCm + in finalCm + +-- private +-- Rebuilds the classmap: merges old class infos into root class infos +-- Returns list of modified root classes +egRebuildClassMap :: (Semigroup d, Eq (f ()), Hashable (f ())) => IntLikeSet EClassId -> State (EGraph d f) ClassReplacements +egRebuildClassMap touchedClasses = state $ \eg -> + let ef = egEquivFind eg + -- Find roots corresponding to all touched classes + roots = ILS.map (`efLookupRoot` ef) touchedClasses + -- Prepare a replacement map for external consumers that just contains changed classes + classReplacements = efSubset (ILS.toList roots) ef + -- Rebuild the class map + cm' = foldl' (\cm (r, vs) -> egRebuildClassSingle r vs cm) (egClassMap eg) (ILM.toList (efFwd classReplacements)) + in (classReplacements, eg { egClassMap = cm' }) + +-- private +-- Rebuilds the 'EGraph' - merges classes as requested in the worklist and recanonicalizes. +-- This may take several rounds as changes propagate "upward" to parents. +-- Returns +-- 1. class remapping (roots -> removed classes) +-- 2. touched root classes +egRebuild :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ())) => WorkList -> State (EGraph d f) (ClassReplacements, IntLikeSet EClassId) +egRebuild wl0 = goRec where + goRec = do + -- Note the existing hashcons + origHc <- gets egHashCons + -- Merge and induce equivalences + -- We track "touched classes" to know which to later rebuild in the classmap + tc <- goNodeRounds origHc ILS.empty wl0 ILS.empty + -- Compute final "touched roots" + ef <- gets egEquivFind + let tr = ILS.fromList [y | x <- ILS.toList tc, y <- maybeToList (efFindRoot x ef)] + -- Now everything is merged, so rewrite the changed parts of the classmap + rm <- egRebuildClassMap tc + -- Finally, cleanup all "dead" classes and nodes + egCompact rm + -- And return the final class remapping and touched roots + pure (rm, tr) + goNodeRounds !origHc !tc !wl !parents = + if null wl && ILS.null parents + then pure tc + else do + (newTc, newWl, newParents) <- egRebuildNodeRound origHc wl parents + let mergedTc = ILS.union newTc tc + goNodeRounds origHc mergedTc newWl newParents + +-- private +-- Replace parent nodes with correct (remapped) ones +egCompactParentClass :: IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f +egCompactParentClass nodeReplacements (EClassInfo dat nodes parents) = + EClassInfo dat nodes (ILS.map (\n -> ILM.findWithDefault n n nodeReplacements) parents) + +-- private +-- Remove dead nodes from given class info +egCompactSelfClass :: (Eq (f ()), Hashable (f ())) => IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f +egCompactSelfClass nodeReplacements (EClassInfo dat nodes parents) = + EClassInfo dat (assocRemoveAllInc (ILM.keys nodeReplacements) nodes) parents + +-- private +-- Find all classes that have dead nodes +findDeadNodeParentClasses :: Foldable f => Assoc ENodeId (f EClassId) -> [ENodeId] -> IntLikeSet EClassId +findDeadNodeParentClasses assoc = foldl' go ILS.empty where + go s n = foldl' (flip ILS.insert) s (assocPartialLookupByKey n assoc) + +-- private +-- Remove all dead nodes and classes from the graph +egCompactInc :: (Foldable f, Eq (f ()), Hashable (f ())) => ClassReplacements -> EGraph d f -> EGraph d f +egCompactInc rm eg = + let ef = egEquivFind eg + assoc = egNodeAssoc eg + hc = egHashCons eg + cm = egClassMap eg + deadClasses = ILM.keysSet (efBwd rm) + -- remove dead nodes from assoc + (nodeReplacements, assoc') = assocCompactInc assoc + -- select all live classes that are parents of dead nodes + deadNodeParentClasses = findDeadNodeParentClasses assoc (ILM.keys nodeReplacements) + -- select all live classes that contain dead nodes + deadNodeSelfClasses = ILS.fromList (fmap (`ILM.partialLookup` hc) (ILM.keys nodeReplacements)) + -- remove dead classes from hashcons + hc' = foldl' (flip ILM.delete) hc (ILM.keys nodeReplacements) + -- remove dead classes from unionfind + (_, ef') = efCompactInc ef + -- remove dead classes from classmap + cm' = foldl' (flip ILM.delete) cm (ILS.toList deadClasses) + -- rewrite dead parent nodes in classmap + cm'' = foldl' (flip (ILM.adjust (egCompactParentClass nodeReplacements))) cm' (ILS.toList deadNodeParentClasses) + -- rewrite dead self nodes in classmap + cm''' = foldl' (flip (ILM.adjust (egCompactSelfClass nodeReplacements))) cm'' (ILS.toList deadNodeSelfClasses) + in eg { egEquivFind = ef', egNodeAssoc = assoc', egClassMap = cm''', egHashCons = hc' } + +-- private +egCompact :: (Foldable f, Eq (f ()), Hashable (f ())) => ClassReplacements -> State (EGraph d f) () +egCompact = modify' . egCompactInc + +-- | Reanalyze a subset of classes - touched roots from merging is sufficient to ensure +-- complete reanalysis. (Note this is implemented in a simplistic way, just taking the +-- fixed point of rounds of analysis. The number of rounds can be no more than the size +-- of the given set.) +-- It may be necessary to call this because merging may leave class analyses in an +-- under-approximating state. This method gives you the true analysis by fixed point. +egReanalyzeSubset :: (Eq d, Semigroup d, Functor f) => EAnalysis d f -> IntLikeSet EClassId -> State (EGraph d f) () +egReanalyzeSubset ana tr = goStart where + goStart = do + cm <- gets egClassMap + let am = ILM.map eciData cm + goRec am + goRec am0 = do + cm <- gets egClassMap + assoc <- gets egNodeAssoc + let fwd = assocFwd assoc + let onNode n = + let fc = ILM.partialLookup n fwd + fd = fmap (`ILM.partialLookup` am0) fc + in ana fd + let calcClass cr = + let nodes = eciNodes (ILM.partialLookup cr cm) + in case ILM.keys (assocFwd nodes) of + n0:ns -> + let d0 = onNode n0 + in foldl' (\d n -> d <> onNode n) d0 ns + [] -> error "impossible" + let onClassRoot p@(_, amx) cr = + let d0 = ILM.partialLookup cr am0 + d1 = calcClass cr + in if d0 == d1 then p else (True, ILM.insert cr d1 amx) + let (changed, am1) = foldl' onClassRoot (False, am0) (ILS.toList tr) + if changed + then goRec am1 + else modify' $ \eg -> + let cm0 = egClassMap eg + cm1 = ILM.mapWithKey (\i c -> c { eciData = ILM.partialLookup i am1 } ) cm0 + in eg { egClassMap = cm1 } + +-- | Reanalyze all classes in the graph. +egReanalyze :: (Eq d, Semigroup d, Functor f) => EAnalysis d f -> State (EGraph d f) () +egReanalyze ana = egClasses >>= egReanalyzeSubset ana diff --git a/src/Overeasy/EquivFind.hs b/src/Overeasy/EquivFind.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/EquivFind.hs @@ -0,0 +1,350 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | See 'EquivFind'. +module Overeasy.EquivFind + ( EquivFind + , efFwd + , efBwd + , efRootsSize + , efLeavesSize + , efTotalSize + , efCanonicalize + , efCanonicalizePartial + , efNew + , efSingleton + , efMember + , efRoots + , efLeaves + , efMembers + , EquivAddRes (..) + , efAddInc + , efAdd + , efEquivs + , efClosure + , efFindRoot + , efFindLeaves + , efSubset + , efLookupRoot + , efLookupLeaves + , efFindAll + , EquivMergeRes (..) + , efUnsafeMerge + , efMergeInc + , efMerge + , EquivMergeSetsRes (..) + , efMergeSetsInc + , efMergeSets + , efCanCompact + , efCompactInc + , efCompact + , efRemoveAllInc + , efRemoveAll + , efUnsafeAddLeafInc + ) where + +import Control.Applicative ((<|>)) +import Control.DeepSeq (NFData) +import Control.Monad.State.Strict (State, state) +import Data.Coerce (Coercible) +import Data.Foldable (foldl') +import Data.Maybe (fromJust, fromMaybe) +import GHC.Generics (Generic) +import IntLike.Map (IntLikeMap) +import qualified IntLike.Map as ILM +import IntLike.Set (IntLikeSet) +import qualified IntLike.Set as ILS + +-- | A "Union-Find" implementation using explicit equivalence classes. +-- Sure, the asympotics aren't as good, but we eventually have to construct these +-- classes, so we might as well just do it as we go! +data EquivFind x = EquivFind + { efFwd :: !(IntLikeMap x (IntLikeSet x)) + -- ^ Map of root to equivalent leaves + -- Invariant: Map keys are only roots + -- Invariant: Sets only contain leaf keys (and not the root itself) + , efBwd :: !(IntLikeMap x x) + -- ^ Map of leaf to root + -- Invariant: Map keys are only leaves, values are only roots + } deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Number of roots in the equiv. +efRootsSize :: EquivFind x -> Int +efRootsSize = ILM.size . efFwd + +-- | Number of leaves in the equiv. +efLeavesSize :: EquivFind x -> Int +efLeavesSize = ILM.size . efBwd + +-- | Total number of keys in the equiv. +efTotalSize :: EquivFind x -> Int +efTotalSize ef = efRootsSize ef + efLeavesSize ef + +-- | Canonicalize the given expression functor by replacing leaves with roots. +-- If any elements are missing, the first is returned. +efCanonicalize :: (Traversable f, Coercible x Int) => f x -> EquivFind x -> Either x (f x) +efCanonicalize fx ef = traverse (\x -> maybe (Left x) pure (efFindRoot x ef)) fx + +-- | Canonicalize the given expression functor by replacing leaves with roots. +-- If any elements are missing, they are simply skipped. +efCanonicalizePartial :: (Functor f, Coercible x Int) => f x -> EquivFind x -> f x +efCanonicalizePartial fx ef = fmap (`efLookupRoot` ef) fx + +-- | Creates an empty equiv +efNew :: EquivFind x +efNew = EquivFind ILM.empty ILM.empty + +-- | Creates a singleton equiv +efSingleton :: Coercible x Int => x -> EquivFind x +efSingleton x = EquivFind (ILM.singleton x ILS.empty) ILM.empty + +-- private +allocMM :: Coercible x Int => x -> IntLikeMap x (IntLikeSet x) -> IntLikeMap x (IntLikeSet x) +allocMM = ILM.alter (<|> Just ILS.empty) + +-- private +insertMM :: Coercible x Int => x -> x -> IntLikeMap x (IntLikeSet x) -> IntLikeMap x (IntLikeSet x) +insertMM x y = ILM.alter (\case { Nothing -> Just (ILS.singleton y); Just s -> Just (ILS.insert y s) }) x + +-- | Result of adding something to the equiv, if you're interested. +data EquivAddRes x = + EquivAddResAlreadyRoot + | EquivAddResAlreadyLeafOf !x + | EquivAddResNewRoot + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Add the given key to the equiv (raw version). +efAddInc :: Coercible x Int => x -> EquivFind x -> (EquivAddRes x, EquivFind x) +efAddInc x ef@(EquivFind fwd bwd) = + case ILM.lookup x bwd of + Nothing -> + if ILM.member x fwd + then (EquivAddResAlreadyRoot, ef) + else (EquivAddResNewRoot, EquivFind (ILM.insert x ILS.empty fwd) bwd) + Just y -> (EquivAddResAlreadyLeafOf y, ef) + +-- | Add the given key to the equiv (raw version). +efAdd :: Coercible x Int => x -> State (EquivFind x) (EquivAddRes x) +efAdd = state . efAddInc + +-- | All keys equivalent to the given key in the equiv. +-- Always returns a set with the given key, even if it's not present. +efEquivs :: Coercible x Int => x -> EquivFind x -> IntLikeSet x +efEquivs x ef = let r = efLookupRoot x ef in ILS.insert r (efLookupLeaves r ef) + +-- | Set of all keys equivalent to the given keys in the equiv. +efClosure :: Coercible x Int => [x] -> EquivFind x -> IntLikeSet x +efClosure xs ef = foldl' (\c x -> if ILS.member x c then c else ILS.union (efEquivs x ef) c) ILS.empty xs + +-- | Find the root equivalent to the given key (if it exists). +efFindRoot :: Coercible x Int => x -> EquivFind x -> Maybe x +efFindRoot x ef = ILM.lookup x (efBwd ef) <|> if ILM.member x (efFwd ef) then Just x else Nothing + +-- | Find the leaves equivalent to the given key (if they exist). +efFindLeaves :: Coercible x Int => x -> EquivFind x -> Maybe (IntLikeSet x) +efFindLeaves x ef = ILM.lookup x (efFwd ef) + +-- | Returns an EquivFind subset representing the given list of keys. +efSubset :: Coercible x Int => [x] -> EquivFind x -> EquivFind x +efSubset xs0 ef0 = foldl' go efNew xs0 where + go (EquivFind fwd1 bwd1) x = + let r = efLookupRoot x ef0 + ls = efLookupLeaves r ef0 + in EquivFind (ILM.insert r ls fwd1) (foldl' (\b l -> ILM.insert l r b) bwd1 (ILS.toList ls)) + +-- | Like 'efFindRoot' but returns same key if not found - does not guarantee presence in map. +efLookupRoot :: Coercible x Int => x -> EquivFind x -> x +efLookupRoot x = fromMaybe x . ILM.lookup x . efBwd + +-- | Like 'efFindLeaves' but returns empty set if not found - does not guarantee presence in map. +efLookupLeaves :: Coercible x Int => x -> EquivFind x -> IntLikeSet x +efLookupLeaves x = fromMaybe ILS.empty . ILM.lookup x . efFwd + +-- | Returns the set of roots for the given set of keys, or an error with the first key +-- not found in the equiv. +efFindAll :: Coercible x Int => [x] -> EquivFind x -> Either x (IntLikeSet x) +efFindAll xs ef = go ILS.empty xs where + go !acc = \case + [] -> Right acc + y:ys -> + case efFindRoot y ef of + Nothing -> Left y + Just z -> go (ILS.insert z acc) ys + +-- | Is the key in the equiv? +efMember :: Coercible x Int => x -> EquivFind x -> Bool +efMember x (EquivFind fwd bwd) = ILM.member x fwd || ILM.member x bwd + +-- | List all roots in the equiv. +efRoots :: Coercible x Int => EquivFind x -> [x] +efRoots = ILM.keys . efFwd + +-- | List all leaves in the equiv. +efLeaves :: Coercible x Int => EquivFind x -> [x] +efLeaves = ILM.keys . efBwd + +-- | List all members (roots and leaves) in the equiv. +efMembers :: Coercible x Int => EquivFind x -> [x] +efMembers ef = efRoots ef ++ efLeaves ef + +-- | The result of trying to merge two keys, if you care. +data EquivMergeRes x = + EquivMergeResMissing !x + | EquivMergeResUnchanged !x + | EquivMergeResChanged !x !(IntLikeSet x) !(EquivFind x) + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Don't even think about it, it's got unsafe in the name. +efUnsafeMerge :: (Coercible x Int, Ord x) => x -> x -> EquivFind x -> (x, IntLikeSet x, EquivFind x) +efUnsafeMerge ix jx (EquivFind fwd bwd) = + let loKey = min ix jx + hiKey = max ix jx + hiSet = ILS.insert hiKey (ILM.partialLookup hiKey fwd) + finalFwd = ILM.adjust (hiSet <>) loKey (ILM.delete hiKey fwd) + finalBwd = foldl' (flip (`ILM.insert` loKey)) bwd (ILS.toList hiSet) + in (loKey, hiSet, EquivFind finalFwd finalBwd) + +-- | Merge two keys (raw version). +efMergeInc :: (Coercible x Int, Ord x) => x -> x -> EquivFind x -> EquivMergeRes x +efMergeInc i j ef = + case efFindRoot i ef of + Nothing -> EquivMergeResMissing i + Just ix -> + case efFindRoot j ef of + Nothing -> EquivMergeResMissing j + Just jx -> + if ix == jx + then EquivMergeResUnchanged ix + else + let (loKey, hiSet, ef') = efUnsafeMerge ix jx ef + in EquivMergeResChanged loKey hiSet ef' + +-- | Merge two keys (state version). +efMerge :: (Coercible x Int, Ord x) => x -> x -> State (EquivFind x) (Maybe (x, IntLikeSet x)) +efMerge i j = state $ \ef -> + case efMergeInc i j ef of + EquivMergeResChanged loKey hiSet ef' -> (Just (loKey, hiSet), ef') + _ -> (Nothing, ef) + +-- | The result of trying to merge multiple keys, if you care. +data EquivMergeManyRes x = + EquivMergeManyResEmpty + | EquivMergeManyResEmbed !(EquivMergeRes x) + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | The result of trying to merge multiple sets of keys, if you care. +data EquivMergeSetsRes x = + EquivMergeSetsResEmptySet + | EquivMergeSetsResMissing !x + | EquivMergeSetsResUnchanged !(IntLikeSet x) + | EquivMergeSetsResChanged !(IntLikeSet x) !(IntLikeSet x) !(EquivFind x) + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Merge sets of keys (raw version). +efMergeSetsInc :: Coercible x Int => [IntLikeSet x] -> EquivFind x -> EquivMergeSetsRes x +efMergeSetsInc css0 u0 = res where + res = + case css0 of + [] -> EquivMergeSetsResUnchanged ILS.empty + _ -> go ILS.empty ILS.empty u0 css0 + go !roots !classRemapSet ef@(EquivFind fwd bwd) css = + case css of + [] -> + let finalRoots = ILS.map (`efLookupRoot` ef) roots + in if ILS.null classRemapSet + then EquivMergeSetsResUnchanged finalRoots + else EquivMergeSetsResChanged finalRoots classRemapSet ef + ds:dss -> + case ILS.toList ds of + [] -> go roots classRemapSet ef dss + zs -> case efFindAll zs ef of + Left x -> EquivMergeSetsResMissing x + Right xs -> + let (loKey, ys) = fromJust (ILS.minView xs) + newRoots = ILS.insert loKey roots + hiSet = ILS.unions (fmap (\y -> ILS.insert y (efLookupLeaves y ef)) (ILS.toList ys)) + newClassRemapSet = ILS.union hiSet classRemapSet + newFwd = ILM.adjust (ILS.union hiSet) loKey (foldl' (flip ILM.delete) fwd (ILS.toList ys)) + newBwd = foldl' (flip (`ILM.insert` loKey)) bwd (ILS.toList hiSet) + newU = EquivFind newFwd newBwd + in go newRoots newClassRemapSet newU dss + +-- | Merge sets of keys (state version). +efMergeSets :: Coercible x Int => [IntLikeSet x] -> State (EquivFind x) (Maybe (IntLikeSet x, IntLikeSet x)) +efMergeSets css = state $ \ef -> + case efMergeSetsInc css ef of + EquivMergeSetsResChanged roots classRemapSet ef' -> (Just (roots, classRemapSet), ef') + _ -> (Nothing, ef) + +-- | Are they compactible keys? +efCanCompact :: EquivFind x -> Bool +efCanCompact = not . ILM.null . efBwd + +-- | See 'efCompact' (this is the raw version). +efCompactInc :: Coercible x Int => EquivFind x -> (IntLikeMap x (IntLikeSet x), EquivFind x) +efCompactInc (EquivFind origFwd origBwd) = finalRes where + finalRes = + let (rootMap, fwd') = foldl' go (ILM.empty, origFwd) (ILM.elems origBwd) + in (rootMap, EquivFind fwd' ILM.empty) + go p@(rootMap, fwd) r = + if ILM.member r rootMap + then p + else + let xs = ILM.partialLookup r fwd + in (ILM.insert r xs rootMap, if ILS.null xs then fwd else ILM.insert r ILS.empty fwd) + +-- | Removes leaves and returns map of root to deleted leaf. +efCompact :: Coercible x Int => State (EquivFind x) (IntLikeMap x (IntLikeSet x)) +efCompact = state efCompactInc + +-- | See 'efRemoveAll' (this is the raw version). +efRemoveAllInc :: Coercible x Int => [x] -> EquivFind x -> (IntLikeMap x x, EquivFind x) +efRemoveAllInc xs (EquivFind fwd0 bwd0) = (remapFinal, EquivFind fwdFinal bwdFinal) where + (fwdFinal, bwdFinal, remapFinal) = foldl' go (fwd0, bwd0, ILM.empty) xs + go tup@(fwd, bwd, remap) x = + case ILM.lookup x fwd of + -- Key is not root + Nothing -> case ILM.lookup x bwd of + -- Key is missing, skip it + Nothing -> tup + -- Key is leaf, remove from both containers + Just r -> + let bwd' = ILM.delete x bwd + fwd' = ILM.adjust (ILS.delete x) r fwd + in (fwd', bwd', remap) + -- Key is root + Just leaves -> + -- ensure the remapping is from ORIGINAL roots to new roots + let origRoot = fromMaybe x (ILM.lookup x bwd0) + in case ILS.minView leaves of + -- Singleton root, remove from fwd and remap + Nothing -> + let fwd' = ILM.delete x fwd + remap' = ILM.delete origRoot remap + in (fwd', bwd, remap') + -- Non-singleton root, rotate + Just (y, rest) -> + let fwd' = ILM.delete x (ILM.insert y rest fwd) + bwd' = ILM.delete y (foldl' (\m l -> ILM.insert l y m) bwd (ILS.toList rest)) + remap' = ILM.insert origRoot y remap + in (fwd', bwd', remap') + +-- | Removes the given keys from the equiv map. +-- If a key is a leaf or singleton root, simply remove it. +-- If it is a root of a larger class, select the min leaf and make it root. +-- Returns a map of old roots to new roots (only those changed in the process - +-- possibly empty). If a key is not found, it is simply ignored. +efRemoveAll :: Coercible x Int => [x] -> State (EquivFind x) (IntLikeMap x x) +efRemoveAll = state . efRemoveAllInc + +-- | Given root, add leaf. Requires that root be present in the map +-- and that leaf would be picked as a leaf. (Therefore, unsafe.) +-- Exposed for efficient merging. +efUnsafeAddLeafInc :: Coercible x Int => x -> x -> EquivFind x -> EquivFind x +efUnsafeAddLeafInc root leaf ef@(EquivFind fwd bwd) = + let trueRoot = efLookupRoot root ef + in EquivFind (ILM.adjust (ILS.insert leaf) trueRoot fwd) (ILM.insert leaf trueRoot bwd) diff --git a/src/Overeasy/Matching.hs b/src/Overeasy/Matching.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/Matching.hs @@ -0,0 +1,348 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Methods to match patterns in 'EGraph's (aka e-matching) +module Overeasy.Matching + ( Pat + , patVars + , Subst + , substVars + , VarId (..) + , PatGraphC + , PatGraph (..) + , patGraph + , Match (..) + , MatchPat (..) + , MatchF (..) + , MatchPatF (..) + , matchVars + , matchClasses + , MatchSubst (..) + , SolGraphC + , SolGraph (..) + , solGraph + , SolStream + , SolveC + , solve + , match + ) where + +import Control.Applicative (Alternative (..)) +import Control.DeepSeq (NFData) +import Control.Monad (void) +import Control.Monad.Reader (asks) +import Control.Monad.State.Strict (MonadState (..), State, evalState, execState, gets, modify', runState) +import Data.Bifunctor (bimap) +import Data.Coerce (Coercible) +import Data.Foldable (fold, foldMap', foldl', for_, toList) +import Data.Functor.Foldable (Base, Corecursive (..), Recursive (..), cata) +import Data.Hashable (Hashable) +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap +import Data.HashSet (HashSet) +import qualified Data.HashSet as HashSet +import IntLike.Map (IntLikeMap) +import qualified IntLike.Map as ILM +import IntLike.Set (IntLikeSet) +import qualified IntLike.Set as ILS +import Overeasy.Assoc (Assoc, assocBwd, assocEquiv, assocFootprint, assocFwd, assocInsertInc, assocLookupByValue, + assocNew) +import Overeasy.EGraph (EClassId (..), EGraph (egHashCons), ENodeId (..), eciNodes, egClassMap, egNodeAssoc) +import Overeasy.EquivFind (efLookupRoot) +import Overeasy.Source (Source, sourceAddInc, sourceNew) +import Overeasy.Streams (Stream, chooseWith, streamAll) +import Unfree (Free, FreeF (..)) + +-- | A pattern is exactly the free monad over the expression functor +-- It has spots for var names ('FreePure') and spots for structural +-- pieces ('FreeEmbed') +type Pat = Free + +-- | The base functor of 'Pat'. +type PatF = FreeF + +-- | The set of vars for a pattern. +patVars :: (Foldable f, Eq v, Hashable v) => Pat f v -> HashSet v +patVars = foldMap' HashSet.singleton + +-- | A substitution. +type Subst c v = HashMap v c + +-- | The set of vars for a substitution. +substVars :: Subst a v -> HashSet v +substVars = HashMap.keysSet + +-- | A match is a pattern annotated with classes (or other data). +data Match c f v = Match + { matchAnno :: !c + , matchPat :: !(MatchPat c f v) + } deriving stock (Functor, Foldable, Traversable) + +deriving stock instance (Eq c, Eq v, Eq (f (Match c f v))) => Eq (Match c f v) +deriving stock instance (Show c, Show v, Show (f (Match c f v))) => Show (Match c f v) + +-- | Tie the knot - the inner layer of a match. +data MatchPat c f v = + MatchPatPure !v + | MatchPatEmbed !(f (Match c f v)) + deriving stock (Functor, Foldable, Traversable) + +deriving stock instance (Eq v, Eq (f (Match c f v))) => Eq (MatchPat c f v) +deriving stock instance (Show v, Show (f (Match c f v))) => Show (MatchPat c f v) + +-- | The base functor of 'Match' +data MatchF c f v r = MatchF + { matchClassF :: !c + , matchPatF :: !(MatchPatF f v r) + } deriving stock (Functor, Foldable, Traversable) + +-- | Tie the knot - the inner part of 'MatchF'. +data MatchPatF f v r = + MatchPatPureF !v + | MatchPatEmbedF !(f r) + deriving stock (Functor, Foldable, Traversable) + +type instance Base (Match c f v) = MatchF c f v + +instance Functor f => Recursive (Match c f v) where + project (Match cl mp) = MatchF cl $ case mp of + MatchPatPure v -> MatchPatPureF v + MatchPatEmbed f -> MatchPatEmbedF f + +instance Functor f => Corecursive (Match c f v) where + embed (MatchF cl mpf) = Match cl $ case mpf of + MatchPatPureF v -> MatchPatPure v + MatchPatEmbedF f -> MatchPatEmbed f + +-- | The set of vars in a match. +matchVars :: (Foldable f, Eq v, Hashable v) => Match c f v -> HashSet v +matchVars = foldMap' HashSet.singleton + +-- | The set of classes in a match. +matchClasses :: (Coercible c Int, Functor f, Foldable f) => Match c f v -> IntLikeSet c +matchClasses = cata go where + go (MatchF cl mpf) = ILS.insert cl $ case mpf of + MatchPatPureF _ -> ILS.empty + MatchPatEmbedF fc -> fold fc + +-- | A apri of match and substitution. +data MatchSubst c f v = MatchSubst + { msMatch :: !(Match c f v) + , msSubst :: !(Subst c v) + } + +deriving stock instance (Eq c, Eq v, Eq (f (Match c f v))) => Eq (MatchSubst c f v) +deriving stock instance (Show c, Show v, Show (f (Match c f v))) => Show (MatchSubst c f v) + +-- | An opaque var id +-- Constructor exported for coercibility +newtype VarId = VarId { unVarId :: Int } + deriving stock (Show) + deriving newtype (Eq, Ord, Enum, Hashable, NFData) + +-- | A pattern graph - can be created once for each pattern and reused +-- for many iterations of search. +data PatGraph f v = PatGraph + { pgRoot :: !VarId + , pgNodes :: !(IntLikeMap VarId (PatF f v VarId)) + , pgVars :: !(HashMap v VarId) + } + +deriving stock instance (Eq v, Eq (f VarId)) => Eq (PatGraph f v) +deriving stock instance (Show v, Show (f VarId)) => Show (PatGraph f v) + +-- | The set of constraints necessary to build a pattern graph. +type PatGraphC f v = (Traversable f, Eq v, Eq (f VarId), Hashable v, Hashable (f VarId)) + +data GraphState f v = GraphState + { gsSrc :: !(Source VarId) + , gsAssoc :: !(Assoc VarId (PatF f v VarId)) + } + +emptyGraphState :: GraphState f v +emptyGraphState = GraphState (sourceNew (VarId 0)) assocNew + +graphEnsurePart :: PatGraphC f v => PatF f v VarId -> State (GraphState f v) VarId +graphEnsurePart part = do + mi <- gets (assocLookupByValue part . gsAssoc) + case mi of + Just i -> pure i + Nothing -> state $ \st -> + let (i, src') = sourceAddInc (gsSrc st) + (_, assoc') = assocInsertInc i part (gsAssoc st) + in (i, st { gsSrc = src', gsAssoc = assoc' }) + +graphEnsurePat :: PatGraphC f v => Pat f v -> State (GraphState f v) VarId +graphEnsurePat = cata go where + go = \case + FreePureF v -> graphEnsurePart (FreePureF v) + FreeEmbedF fp -> do + fi <- sequenceA fp + graphEnsurePart (FreeEmbedF fi) + +graphCanonicalize :: PatGraphC f v => GraphState f v -> IntLikeMap VarId (PatF f v VarId) +graphCanonicalize (GraphState _ assoc) = + let fwd = assocFwd assoc + equiv = assocEquiv assoc + in fmap (fmap (`efLookupRoot` equiv)) fwd + +-- | Builds a pattern graph from a pattern. +patGraph :: PatGraphC f v => Pat f v -> PatGraph f v +patGraph p = + let (i, st) = runState (graphEnsurePat p) emptyGraphState + m = graphCanonicalize st + n = HashMap.fromList (ILM.toList m >>= \(j, x) -> case x of { FreePureF v -> [(v, j)]; _ -> [] }) + in PatGraph i m n + +-- | A solution graph - must be created from an e-graph each merge/rebuild. +data SolGraph c f = SolGraph + { sgByVar :: !(IntLikeMap VarId (IntLikeSet c)) + -- ^ Map of var -> classes. + -- Contains all vars. + -- If the inner map is empty, that means the pattern was not matched. + -- The inner set will not be empty. + , sgNodes :: !(HashMap (f c) c) + -- ^ Map of node structures to classes. + } + +deriving stock instance (Eq c, Eq (f c)) => Eq (SolGraph c f) +deriving stock instance (Show c, Show (f c)) => Show (SolGraph c f) + +-- | The set of constraints necessary to build a solution graph. +type SolGraphC f = (Functor f, Foldable f, Eq (f ()), Hashable (f ())) + +-- | Builds a solution graph from an e-graph. +solGraph :: SolGraphC f => PatGraph f v -> EGraph d f -> SolGraph EClassId f +solGraph pg eg = + -- For each class, use footprint of reverse node assoc to find set of node ids + -- Start with just the embedded nodes + let byVarEmbed = ILM.fromList $ ILM.toList (pgNodes pg) >>= \(i, pf) -> + case pf of + FreePureF _ -> [] + FreeEmbedF fi -> + let fu = void fi + cns = ILM.toList (egClassMap eg) >>= \(c, inf) -> + let ns = eciNodes inf + fp = assocFootprint fu ns + in [(c, fp) | not (ILS.null fp)] + in [(i, bimap ILS.fromList mconcat (unzip cns))] + byVar = genByVar byVarEmbed (pgNodes pg) (assocFwd (egNodeAssoc eg)) + hc = egHashCons eg + nodes = fmap (`ILM.partialLookup` hc) (assocBwd (egNodeAssoc eg)) + in SolGraph byVar nodes + +data Record = + RecordPure !VarId !(IntLikeSet EClassId) + | RecordEmbed + deriving stock (Eq, Show) + +type Records = [Record] + +initRecords :: Foldable f => IntLikeMap VarId (PatF f v VarId) -> f VarId -> Records +initRecords nodes = fmap (\i -> case ILM.partialLookup i nodes of { FreePureF _ -> RecordPure i ILS.empty; _ -> RecordEmbed }) . toList + +updateRecords :: Foldable f => Records -> f EClassId -> Records +updateRecords rs = zipWith (\r c -> case r of { RecordPure v cs -> RecordPure v (ILS.insert c cs); _ -> r } ) rs . toList + +genByVar :: Foldable f => IntLikeMap VarId (IntLikeSet EClassId, IntLikeSet ENodeId) -> IntLikeMap VarId (PatF f v VarId) -> IntLikeMap ENodeId (f EClassId) -> IntLikeMap VarId (IntLikeSet EClassId) +genByVar byVarEmbed nodes fwd = execState (for_ (ILM.toList nodes) go) (fmap fst byVarEmbed) where + go (i, pf) = + case pf of + FreePureF _ -> pure () + FreeEmbedF fi -> do + -- We've gone through embedded patterns before so we're able + -- to look up nodes for each pattern + let (_, ns) = ILM.partialLookup i byVarEmbed + -- For each node, update positionally what it could be + let rs = foldl' (\rsx n -> let fc = ILM.partialLookup n fwd in updateRecords rsx fc) (initRecords nodes fi) (ILS.toList ns) + -- Finally update the map; if missing set the positions as is, otherwise take intersection + modify' $ \m -> foldl' (\mx r -> case r of {RecordPure j cs -> ILM.alter (Just . maybe cs (ILS.intersection cs)) j mx; _ -> mx}) m rs + +data SolEnv c f v = SolEnv + { sePatGraph :: !(PatGraph f v) + , seSolGraph :: !(SolGraph c f) + } + +newtype SolState c = SolState + { ssClasses :: IntLikeMap VarId c + } deriving (Eq, Show) + +-- | A stream of solutions. Can be demanded all at once, unconsed one at a time, +-- or interleaved. +type SolStream c f v z = Stream (SolEnv c f v) (SolState c) z + +-- | The set of constraints necessary to search for solutions. +type SolveC c f v = (Traversable f, Coercible c Int, Eq v, Hashable v, Eq (f c), Hashable (f c)) + +constructMatch :: Traversable f => IntLikeMap VarId (PatF f v VarId) -> IntLikeMap VarId c -> VarId -> Match c f v +constructMatch nodes classes i0 = evalState (go i0) ILM.empty where + go i = do + cache <- get + case ILM.lookup i cache of + Just res -> pure res + Nothing -> do + let c = ILM.partialLookup i classes + mp <- case ILM.partialLookup i nodes of + FreePureF v -> pure $! MatchPatPure v + FreeEmbedF f -> fmap MatchPatEmbed (traverse go f) + pure $! Match c mp + +constructSubst :: HashMap v VarId -> IntLikeMap VarId a -> Subst a v +constructSubst vars classes = fmap (`ILM.partialLookup` classes) vars + +solveYield :: Traversable f => SolStream c f v (MatchSubst c f v) +solveYield = do + pg <- asks sePatGraph + classes <- gets ssClasses + let mat = constructMatch (pgNodes pg) classes (pgRoot pg) + subst = constructSubst (pgVars pg) classes + ms = MatchSubst mat subst + pure ms + +-- | Produces a stream of solutions (e-matches). +solve :: SolveC c f v => SolStream c f v (MatchSubst c f v) +solve = do + i <- asks (pgRoot . sePatGraph) + void (solveRec i) + solveYield + +solveChoose :: SolveC c f v => VarId -> IntLikeSet c -> SolStream c f v c +solveChoose i cs = chooseWith (ILS.toList cs) (solveSet i) + +solveSet :: VarId -> c -> SolStream c f v c +solveSet i c = + c <$ modify' (\ss -> ss { ssClasses = ILM.insert i c (ssClasses ss) }) + +solveRec :: SolveC c f v => VarId -> SolStream c f v c +solveRec i = do + ms <- gets (ILM.lookup i . ssClasses) + case ms of + -- Seen before, return solution + Just s -> pure s + -- Unseen + Nothing -> do + n <- asks (ILM.partialLookup i . pgNodes . sePatGraph) + case n of + -- Free var, choose a solution for each class in `sgByVar i` + FreePureF _ -> do + cs <- asks (ILM.partialLookup i . sgByVar . seSolGraph) + solveChoose i cs + -- Embedded functor, traverse and emit solution if present + FreeEmbedF fi -> do + fa <- traverse solveRec fi + mc <- asks (HashMap.lookup fa . sgNodes . seSolGraph) + case mc of + Nothing -> empty + Just c -> solveSet i c + +-- | The easiest way to do e-matching: given a pattern and an e-graph, yield the list of matches. +-- Note that it might be more efficient to keep a 'PatGraph' on hand and uncons the matches +-- one by one. +match :: (PatGraphC f v, SolGraphC f, SolveC EClassId f v) => Pat f v -> EGraph d f -> [MatchSubst EClassId f v] +match p eg = + let pg = patGraph p + sg = solGraph pg eg + in if any ILS.null (ILM.elems (sgByVar sg)) + -- If any var id has no patches, the pattern won't match, so don't try to solve + then [] + else streamAll solve (SolEnv pg sg) (SolState ILM.empty) diff --git a/src/Overeasy/Source.hs b/src/Overeasy/Source.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/Source.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | See 'Source'. +module Overeasy.Source + ( Source + , sourceSize + , sourceNew + , sourceAddInc + , sourceAdd + , sourceSkipInc + , sourceSkip + , sourcePeek + ) where + +import Control.DeepSeq (NFData) +import Control.Monad.State.Strict (State, modify', state) +import Data.Coerce (Coercible, coerce) +import GHC.Generics (Generic) + +-- | A source of unique ids +data Source x = Source + { sourceSize :: !Int + -- ^ How many ids have ever been created? + , sourceNextId :: !Int + } deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Creates a new 'Source' from a starting id +sourceNew :: Coercible x Int => x -> Source x +sourceNew = Source 0 . coerce + +-- | Generates the next id from the source (purely) +sourceAddInc :: Coercible x Int => Source x -> (x, Source x) +sourceAddInc (Source s x) = (coerce x, Source (s + 1) (x + 1)) + +-- | Generates the next id from the source (statefully) +sourceAdd :: Coercible x Int => State (Source x) x +sourceAdd = state sourceAddInc + +-- | Skips past the given id (purely) +sourceSkipInc :: Coercible x Int => x -> Source x -> Source x +sourceSkipInc y (Source s x) = + let z = coerce y + in Source (s + 1) (max x (z + 1)) + +-- | Skips past the given id (statefully) +sourceSkip :: Coercible x Int => x -> State (Source x) () +sourceSkip = modify' . sourceSkipInc + +-- | Peek at the next id +sourcePeek :: Coercible x Int => Source x -> x +sourcePeek = coerce . sourceNextId diff --git a/src/Overeasy/Streams.hs b/src/Overeasy/Streams.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/Streams.hs @@ -0,0 +1,40 @@ +-- | Stuff for streaming search results. +module Overeasy.Streams + ( chooseWith + , choose + , Stream + , streamAll + ) where + +import Control.Applicative (Alternative (..)) +import Control.Monad (MonadPlus) +import Control.Monad.Logic (LogicT, MonadLogic, observeAllT) +import Control.Monad.Reader (MonadReader (..), ReaderT (..)) +import Control.Monad.State.Strict (MonadState (..), State, runState) + +newtype M r s a = M { unM :: ReaderT r (State s) a } + deriving newtype ( + Functor, Applicative, Monad, + MonadReader r, MonadState s) + +runM :: M r s a -> r -> s -> (a, s) +runM m r = runState (runReaderT (unM m) r) + +-- | Choose one of many alteratives and process it with the given function. +chooseWith :: (Foldable f, Alternative m) => f a -> (a -> m b) -> m b +chooseWith fa f = foldr ((<|>) . f) empty fa + +-- | Choose one of many alteratives. +choose :: (Foldable f, Alternative m) => f a -> m a +choose fa = chooseWith fa pure + +-- | A stream of results. Just a wrapper around 'LogicT' to keep things tidy. +newtype Stream r s a = Stream { unStream :: LogicT (M r s) a } + deriving newtype ( + Functor, Applicative, Monad, + MonadReader r, MonadState s, MonadLogic, + Alternative, MonadPlus, Semigroup, Monoid) + +-- | Produces all results from the stream. +streamAll :: Stream r s a -> r -> s -> [a] +streamAll stream env st = fst (runM (observeAllT (unStream stream)) env st) diff --git a/src/Overeasy/Util.hs b/src/Overeasy/Util.hs new file mode 100644 --- /dev/null +++ b/src/Overeasy/Util.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | A grab bag of fun stuff. +module Overeasy.Util + ( Whole + , RecursiveWhole + , foldWholeM + , Changed (..) + , stateFail + , stateOption + , stateFailChanged + , stateFold + ) where + +import Control.DeepSeq (NFData) +import Control.Monad (foldM, forM_) +import Control.Monad.State.Strict (State, get, put) +import Data.Functor.Foldable (Base, Recursive (..)) +import Data.Hashable (Hashable) +import GHC.Generics (Generic) + +-- | Often 'f' is primary, not 't'. Relate them with this constraint. +type Whole t f = (f ~ Base t) + +-- | Constraint for recursive structures +type RecursiveWhole t f = (Recursive t, Whole t f) + +-- | Traverses a recursive structure +foldWholeM :: (RecursiveWhole t f, Traversable f, Monad m) => (f a -> m a) -> t -> m a +foldWholeM h = go where + go t = do + let ft = project t + fa <- traverse go ft + h fa + +-- | A nicely-named 'Bool' for tracking state changes +data Changed = ChangedNo | ChangedYes + deriving stock (Eq, Ord, Show, Generic) + deriving anyclass (Hashable, NFData) + +instance Semigroup Changed where + c1 <> c2 = + case c1 of + ChangedYes -> ChangedYes + _ -> c2 + +instance Monoid Changed where + mempty = ChangedNo + mappend = (<>) + +-- | Embeds a function that may fail in a stateful context +stateFail :: (s -> Maybe (b, s)) -> State s (Maybe b) +stateFail f = do + s <- get + case f s of + Nothing -> pure Nothing + Just (b, s') -> put s' >> pure (Just b) + +-- | Embeds a function that may fail in a stateful context +stateOption :: (s -> (b, Maybe s)) -> State s b +stateOption f = do + s <- get + let (b, ms) = f s + forM_ ms put + pure b + +-- | Embeds a function that may fail in a stateful context with change tracking +stateFailChanged :: (s -> Maybe s) -> State s Changed +stateFailChanged f = do + s <- get + case f s of + Nothing -> pure ChangedNo + Just s' -> put s' >> pure ChangedYes + +-- -- | Embeds a stateful action in a larger context +-- stateLens :: Lens' s a -> State a b -> State s b +-- stateLens l act = state $ \s -> +-- let (b, a') = runState act (view l s) +-- s' = set l a' s +-- in (b, s') + +-- | 'foldM' specialized and flipped. +stateFold :: Foldable t => b -> t a -> (b -> a -> State s b) -> State s b +stateFold b as f = foldM f b as +{-# INLINE stateFold #-} diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,953 @@ +module Main (main) where + +import Control.DeepSeq (NFData, force) +import Control.Exception (evaluate) +import Control.Monad (foldM, unless, when) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Control.Monad.State.Strict (MonadState (..), State, StateT, evalState, evalStateT, execState, execStateT, gets, + runState) +import Control.Monad.Trans (MonadTrans (..)) +import Data.Bifunctor (bimap) +import Data.Char (chr, ord) +import Data.Coerce (coerce) +import Data.Foldable (for_) +import Data.Hashable (Hashable) +import qualified Data.HashMap.Strict as HashMap +import Data.List (delete) +import Data.Maybe (fromJust, isJust) +import Data.Semigroup (Max (..)) +import qualified Data.Sequence as Seq +import Data.Traversable (for) +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range +import qualified IntLike.Equiv as ILE +import qualified IntLike.Graph as ILG +import IntLike.Map (IntLikeMap) +import qualified IntLike.Map as ILM +import IntLike.Set (IntLikeSet) +import qualified IntLike.Set as ILS +import Overeasy.Assoc (Assoc, AssocInsertRes (..), assocBwd, assocCanCompact, assocCompact, assocEquiv, assocFromList, + assocFwd, assocInsert, assocLeaves, assocMember, assocMembers, assocNew, assocPartialLookupByKey, + assocRoots, assocSize) +import Overeasy.EGraph (EAnalysis, EClassId (..), EClassInfo (..), EGraph (..), ENodeId (..), MergeResult (..), + egAddTerm, egCanonicalize, egClassSize, egFindTerm, egMerge, egMergeMany, egNew, egNodeSize, + noAnalysis) +import Overeasy.EquivFind (EquivFind (..), efAdd, efCanCompact, efCompact, efFindRoot, efLeaves, efLeavesSize, efMember, + efMembers, efMerge, efMergeSets, efNew, efRemoveAll, efRoots, efRootsSize, efTotalSize) +import Overeasy.Matching (Match (..), MatchPat (..), MatchSubst (..), Pat, match) +import Overeasy.Util (Changed (..)) +import PropUnit (DependencyType (..), Gen, MonadTest, PropertyT, Range, TestLimit, TestTree, after, assert, forAll, + testGroup, testMain, testProp, testUnit, (/==), (===)) +import Test.Overeasy.Arith (Arith (..), ArithF (ArithPlusF)) +import Test.Overeasy.BinTree (BinTree, pattern BinTreeBranch, BinTreeF (..), pattern BinTreeLeaf) +import Unfree (pattern FreeEmbed, pattern FreePure) + +fullyEvaluate :: (MonadIO m, NFData a) => a -> m a +fullyEvaluate = liftIO . evaluate . force + +applyS :: Monad m => State s a -> StateT s m a +applyS = state . runState + +testS :: Monad m => (s -> m a) -> StateT s m a +testS p = get >>= lift . p + +applyTestS :: Monad m => State s a -> (a -> s -> m b) -> StateT s m b +applyTestS act check = do + a <- applyS act + s <- get + lift (check a s) + +foldS_ :: (Monad m, Foldable t) => s -> t a -> (a -> StateT s m ()) -> m s +foldS_ z as f = execStateT (for_ as f) z + +runS :: Monad m => s -> StateT s m () -> m () +runS = flip evalStateT + +flipFoldM :: Monad m => b -> [a] -> (b -> a -> m b) -> m b +flipFoldM b as f = foldM f b as + +newtype V = V { unV :: Int } + deriving newtype (Eq, Ord, Hashable, NFData) + +instance Show V where + show = show . fromV + +toV :: Char -> V +toV = V . ord + +fromV :: V -> Char +fromV = chr . unV + +setV :: String -> IntLikeSet V +setV = ILS.fromList . fmap toV + +mapV :: [(Char, Char)] -> IntLikeMap V V +mapV = ILM.fromList . fmap (bimap toV toV) + +multiMapV :: [(Char, String)] -> IntLikeMap V (IntLikeSet V) +multiMapV = ILM.fromList . fmap (bimap toV setV) + +type EF = EquivFind V + +testEfSimple :: TestTree +testEfSimple = testUnit "EF simple" $ runS efNew $ do + testS $ \ef -> do + efRootsSize ef === 0 + efLeavesSize ef === 0 + efTotalSize ef === 0 + efRoots ef === [] + efLeaves ef === [] + efMembers ef === [] + efMember (toV 'a') ef === False + efMember (toV 'c') ef === False + efFwd ef === ILM.empty + efBwd ef === ILM.empty + _ <- applyS (efAdd (toV 'a')) + testS $ \ef -> do + efRootsSize ef === 1 + efLeavesSize ef === 0 + efTotalSize ef === 1 + ILS.fromList (efRoots ef) === setV "a" + ILS.fromList (efLeaves ef) === ILS.empty + efFwd ef === multiMapV [('a', "")] + efBwd ef === ILM.empty + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + testS $ \ef -> do + efRootsSize ef === 3 + efLeavesSize ef === 0 + efTotalSize ef === 3 + ILS.fromList (efRoots ef) === setV "abc" + ILS.fromList (efLeaves ef) === ILS.empty + efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "")] + efBwd ef === ILM.empty + applyTestS (efMerge (toV 'a') (toV 'c')) $ \res ef -> do + res === Just (toV 'a', setV "c") + efRootsSize ef === 2 + efLeavesSize ef === 1 + efTotalSize ef === 3 + ILS.fromList (efRoots ef) === setV "ab" + ILS.fromList (efLeaves ef) === setV "c" + efFwd ef === multiMapV [('a', "c"), ('b', "")] + efBwd ef === mapV [('c', 'a')] + efMembers ef === fmap toV ['a', 'b', 'c'] + efMember (toV 'a') ef === True + efMember (toV 'c') ef === True + applyTestS (efMerge (toV 'c') (toV 'a')) $ \res _ -> res === Nothing + applyTestS (efMerge (toV 'b') (toV 'z')) $ \res _ -> res === Nothing + +resetEf :: StateT EF (PropertyT IO) () +resetEf = do + put efNew + _ <- applyS (efAdd (toV 'a')) + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + _ <- applyS (efMerge (toV 'a') (toV 'c')) + ef <- get + efFwd ef === multiMapV [('a', "c"), ('b', "")] + efBwd ef === mapV [('c', 'a')] + +addExtraEf :: StateT EF (PropertyT IO) () +addExtraEf = do + _ <- applyS (efAdd (toV 'd')) + _ <- applyS (efMerge (toV 'a') (toV 'd')) + ef <- get + efFwd ef === multiMapV [('a', "cd"), ('b', "")] + efBwd ef === mapV [('c', 'a'), ('d', 'a')] + +testEfRemove :: TestTree +testEfRemove = testUnit "EF remove" $ runS efNew $ do + -- remove leav + resetEf + applyTestS (efRemoveAll [toV 'c']) $ \res ef -> do + res === ILM.empty + efFwd ef === multiMapV [('a', ""), ('b', "")] + efBwd ef === mapV [] + -- remove singleton root + resetEf + applyTestS (efRemoveAll [toV 'b']) $ \res ef -> do + res === ILM.empty + efFwd ef === multiMapV [('a', "c")] + efBwd ef === mapV [('c', 'a')] + -- remove non-singleton root + resetEf + applyTestS (efRemoveAll [toV 'a']) $ \res ef -> do + res === mapV [('a', 'c')] + efFwd ef === multiMapV [('b', ""), ('c', "")] + efBwd ef === mapV [] + -- remove all in class (root -> leaf order) + resetEf + applyTestS (efRemoveAll [toV 'a', toV 'c']) $ \res ef -> do + res === ILM.empty + efFwd ef === multiMapV [('b', "")] + efBwd ef === mapV [] + -- remove all in class (leaf -> root order) + resetEf + applyTestS (efRemoveAll [toV 'c', toV 'a']) $ \res ef -> do + res === ILM.empty + efFwd ef === multiMapV [('b', "")] + efBwd ef === mapV [] + -- remove with rotation and leaf + resetEf + addExtraEf + applyTestS (efRemoveAll [toV 'a']) $ \res ef -> do + res === mapV [('a', 'c')] + efFwd ef === multiMapV [('c', "d"), ('b', "")] + efBwd ef === mapV [('d', 'c')] + -- remove with two rotations + resetEf + addExtraEf + applyTestS (efRemoveAll [toV 'a', toV 'c']) $ \res ef -> do + res === mapV [('a', 'd')] + efFwd ef === multiMapV [('d', ""), ('b', "")] + efBwd ef === mapV [] + -- remove with (leaf, rotation) + resetEf + addExtraEf + applyTestS (efRemoveAll [toV 'c', toV 'a']) $ \res ef -> do + res === mapV [('a', 'd')] + efFwd ef === multiMapV [('d', ""), ('b', "")] + efBwd ef === mapV [] + +testEfRec :: TestTree +testEfRec = testUnit "EF rec" $ runS efNew $ do + _ <- applyS (efAdd (toV 'a')) + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + applyTestS (efMerge (toV 'b') (toV 'c')) $ \res ef -> do + res === Just (toV 'b', setV "c") + efRootsSize ef === 2 + efLeavesSize ef === 1 + efTotalSize ef === 3 + ILS.fromList (efRoots ef) === setV "ab" + ILS.fromList (efLeaves ef) === setV "c" + efFwd ef === multiMapV [('a', ""), ('b', "c")] + efBwd ef === mapV [('c', 'b')] + applyTestS (efMerge (toV 'a') (toV 'c')) $ \res ef -> do + res === Just (toV 'a', setV "bc") + efRootsSize ef === 1 + efLeavesSize ef === 2 + efTotalSize ef === 3 + ILS.fromList (efRoots ef) === setV "a" + ILS.fromList (efLeaves ef) === setV "bc" + efFwd ef === multiMapV [('a', "bc")] + efBwd ef === mapV [('b', 'a'), ('c', 'a')] + +testEfMany :: TestTree +testEfMany = testUnit "EF many" $ runS efNew $ do + _ <- applyS (efAdd (toV 'a')) + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + _ <- applyS (efAdd (toV 'd')) + _ <- applyS (efAdd (toV 'e')) + applyTestS (efMergeSets [setV "cde"]) $ \res ef -> do + res === Just (setV "c", setV "de") + efRootsSize ef === 3 + efLeavesSize ef === 2 + efTotalSize ef === 5 + ILS.fromList (efRoots ef) === setV "abc" + ILS.fromList (efLeaves ef) === setV "de" + efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "de")] + efBwd ef === mapV [('d', 'c'), ('e', 'c')] + applyTestS (efMergeSets [setV "abd"]) $ \res ef -> do + res === Just (setV "a", setV "bcde") + efRootsSize ef === 1 + efLeavesSize ef === 4 + efTotalSize ef === 5 + ILS.fromList (efRoots ef) === setV "a" + ILS.fromList (efLeaves ef) === setV "bcde" + efFwd ef === multiMapV [('a', "bcde")] + efBwd ef === mapV [('b', 'a'), ('c', 'a'), ('d', 'a'), ('e', 'a')] + +testEfSets :: TestTree +testEfSets = testUnit "EF sets" $ runS efNew $ do + _ <- applyS (efAdd (toV 'a')) + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + _ <- applyS (efAdd (toV 'd')) + _ <- applyS (efAdd (toV 'e')) + applyTestS (efMergeSets [setV "cde", setV "abc"]) $ \res ef -> do + res === Just (setV "a", setV "bcde") + efRootsSize ef === 1 + efLeavesSize ef === 4 + efTotalSize ef === 5 + ILS.fromList (efRoots ef) === setV "a" + efFwd ef === multiMapV [('a', "bcde")] + +testEfCompact :: TestTree +testEfCompact = testUnit "EF compact" $ runS efNew $ do + _ <- applyS (efAdd (toV 'a')) + _ <- applyS (efAdd (toV 'b')) + _ <- applyS (efAdd (toV 'c')) + _ <- applyS (efAdd (toV 'd')) + _ <- applyS (efAdd (toV 'e')) + testS $ \ef -> assert (not (efCanCompact ef)) + applyTestS (efMergeSets [setV "cde"]) $ \res ef -> do + res === Just (setV "c", setV "de") + efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "de")] + efBwd ef === mapV [('d', 'c'), ('e', 'c')] + assert (efCanCompact ef) + applyTestS efCompact $ \res ef -> do + efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "")] + efBwd ef === ILM.empty + res === multiMapV [('c', "de")] + assert (not (efCanCompact ef)) + +testEfUnit :: TestTree +testEfUnit = testGroup "EF unit" [testEfSimple, testEfRec, testEfMany, testEfSets, testEfCompact, testEfRemove] + +genDistinctPairFromList :: Eq a => [a] -> Gen (a, a) +genDistinctPairFromList = \case + xs@(_:_:_) -> do + a <- Gen.element xs + b <- Gen.element (delete a xs) + pure (a, b) + _ -> error "List needs more than two elements" + +genListOfDistinctPairs :: Eq a => Range Int -> [a] -> Gen [(a, a)] +genListOfDistinctPairs nOpsRange vs = + if length vs < 2 + then pure [] + else Gen.list nOpsRange (genDistinctPairFromList vs) + +genV :: Int -> Gen V +genV maxElems = + let minVal = ord 'a' + maxVal = minVal + maxElems - 1 + in fmap V (Gen.int (Range.linear minVal maxVal)) + +genMembers :: Int -> Gen [V] +genMembers maxElems = do + let nElemsRange = Range.linear 0 maxElems + minVal = ord 'a' + n <- Gen.int nElemsRange + pure (fmap (\i -> V (minVal + i)) [0..n-1]) + +mkInitEf :: [V] -> EF +mkInitEf vs = execState (for_ vs efAdd) efNew + +mkPairsMergedEf :: [(V, V)] -> EF -> EF +mkPairsMergedEf vvs = execState (for_ vvs (uncurry efMerge)) + +mkSetsMergedEf :: [(V, V)] -> EF -> EF +mkSetsMergedEf vvs = execState (for_ (fmap (\(x, y) -> [ILS.fromList [x, y]]) vvs) efMergeSets) + +mkSingleMergedEf :: [(V, V)] -> EF -> EF +mkSingleMergedEf vvs = execState (efMergeSets (fmap (\(x, y) -> ILS.fromList [x, y]) vvs)) + +data MergeStrat = MergeStratPairs | MergeStratSets | MergeStratSingle + deriving stock (Eq, Show, Enum, Bounded) + +genMergeStrat :: Gen MergeStrat +genMergeStrat = Gen.enumBounded + +testEfProp :: TestLimit -> TestTree +testEfProp lim = after AllSucceed "EF unit" $ testProp "EF prop" lim $ do + let maxElems = 50 + -- generate elements + memberList <- forAll (genMembers maxElems) + let memberSet = ILS.fromList memberList + nMembers = ILS.size memberSet + allPairs = ILS.unorderedPairs memberSet + nOpsRange = Range.linear 0 (nMembers * nMembers) + let initEf = mkInitEf memberList + -- assert that sizes indicate nothing is merged + efRootsSize initEf === nMembers + efLeavesSize initEf === 0 + efTotalSize initEf === nMembers + -- assert that find indicates nothing is merged + for_ allPairs $ \(a, b) -> flip evalStateT initEf $ do + x <- applyS (gets (efFindRoot a)) + y <- applyS (gets (efFindRoot b)) + assert (isJust x) + assert (isJust y) + x /== y + -- generate some pairs and merge them + mergePairs <- forAll (genListOfDistinctPairs nOpsRange memberList) + mergeStrat <- forAll genMergeStrat + let mergedUf = + case mergeStrat of + MergeStratPairs -> mkPairsMergedEf mergePairs initEf + MergeStratSets -> mkSetsMergedEf mergePairs initEf + MergeStratSingle -> mkSingleMergedEf mergePairs initEf + -- assert that total size is unchanged + efTotalSize mergedUf === nMembers + -- calculate components by graph reachability + let components = ILG.undirectedComponents mergePairs + -- assert that elements are equal or not according to component + _ <- foldS_ mergedUf allPairs $ \(a, b) -> do + x <- applyS (gets (efFindRoot a)) + y <- applyS (gets (efFindRoot b)) + let aComponent = ILE.lookupClass a components + bComponent = ILE.lookupClass b components + if isJust aComponent && aComponent == bComponent + then x === y + else x /== y + pure () + +type AV = Assoc ENodeId V + +-- | Asserts assoc is compact - should also check 'assertAssocInvariants' +assertAssocCompact :: (MonadTest m, Eq a, Hashable a, Show a) => Assoc ENodeId a -> m () +assertAssocCompact av = do + let fwd = assocFwd av + bwd = assocBwd av + -- Assert that the assoc has been rebuilt + assert $ not (assocCanCompact av) + -- Look at sizes to confirm that assoc could map 1-1 + ILM.size fwd === HashMap.size bwd + -- Go through keys forward + for_ (ILM.toList fwd) $ \(x, fc) -> do + -- Assert is found in backward map AND maps back + HashMap.lookup fc bwd === Just x + -- Go through keys backward + for_ (HashMap.toList bwd) $ \(fc, x) -> + -- Assert is present in forward map AND maps back + ILM.lookup x fwd === Just fc + +-- | Asserts assoc is correctly structured (compact or not) +assertAssocInvariants :: (MonadTest m, Eq a, Hashable a) => Assoc ENodeId a -> m () +assertAssocInvariants av = do + let fwd = assocFwd av + bwd = assocBwd av + equiv = assocEquiv av + -- First check that fwd and bwd are 1-1 + -- Go through keys forward + for_ (ILM.toList fwd) $ \(_, fc) -> do + -- Assert is found in backward map + assert $ HashMap.member fc bwd + -- Go through keys backward + for_ (HashMap.toList bwd) $ \(_, x) -> + -- Assert is present in forward map + assert $ ILM.member x fwd + -- Assert that fwd keys are exactly the equiv roots + ILS.fromList (ILM.keys fwd) === ILS.fromList (efRoots equiv) + +data AssocCase = AssocCase !String ![(Int, Char)] ![(Int, Char, Int, AssocInsertRes Int)] ![(Int, Char)] + +allAssocCases :: [AssocCase] +allAssocCases = + let start = [(0, 'a'), (1, 'b'), (2, 'c')] + in [ AssocCase "base" start [] start + , AssocCase "ident" start + [(0, 'a', 0, AssocInsertResUnchanged)] + start + , AssocCase "superfluous" start + [(4, 'a', 0, AssocInsertResMerged (ILS.singleton 4))] + start + , AssocCase "internal" start + [(0, 'b', 0, AssocInsertResMerged (ILS.singleton 1))] + [(0, 'b'), (2, 'c')] + , AssocCase "external" start + [(0, 'd', 0, AssocInsertResUpdated)] + [(0, 'd'), (1, 'b'), (2, 'c')] + , AssocCase "additional" start + [(4, 'd', 4, AssocInsertResCreated)] + [(0, 'a'), (1, 'b'), (2, 'c'), (4, 'd')] + , AssocCase "chain fwd" start + -- The singleton set in the second result is just the children (and self) of the clobbered node + -- We don't have to lookup the old clobbered nodes for 1 bc when this is used everything will be merged + [(0, 'b', 0, AssocInsertResMerged (ILS.singleton 1)), (1, 'c', 0, AssocInsertResMerged (ILS.singleton 2))] + [(0, 'c')] + , AssocCase "chain bwd" start + -- The set in the second result is not a singleton here because it already had children + [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (0, 'c', 0, AssocInsertResMerged (ILS.fromList [1,2]))] + [(0, 'c')] + , AssocCase "chain self" start + [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'c', 1, AssocInsertResUnchanged)] + [(0, 'a'), (1, 'c')] + , AssocCase "chain change" start + [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'd', 1, AssocInsertResUpdated)] + [(0, 'a'), (1, 'd')] + , AssocCase "chain back id" start + [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (1, 'b', 1, AssocInsertResUpdated)] + [(0, 'a'), (1, 'b')] + , AssocCase "chain back del" start + [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'b', 1, AssocInsertResUpdated)] + [(0, 'a'), (1, 'b')] + , AssocCase "chain change rev" start + [(2, 'd', 2, AssocInsertResUpdated), (1, 'c', 1, AssocInsertResUpdated)] + [(0, 'a'), (1, 'c'), (2, 'd')] + ] + +mkAssoc :: [(Int, Char)] -> AV +mkAssoc rawPairs = + let pairs = fmap (bimap ENodeId toV) rawPairs + in assocFromList pairs + +runAV :: Monad m => [(Int, Char)] -> StateT AV m () -> m () +runAV = runS . mkAssoc + +testAssocCase :: AssocCase -> TestTree +testAssocCase (AssocCase name start act end) = testUnit name $ runAV start $ do + testS $ \av -> do + assertAssocInvariants av + assertAssocCompact av + assocSize av === length start + for_ act $ \(x, a, expectedY, expectedRes) -> do + (actualY, actualRes) <- applyS (assocInsert (ENodeId x) (toV a)) + (actualY, actualRes) === coerce (expectedY, expectedRes) + testS assertAssocInvariants + _ <- applyS assocCompact + testS $ \av -> do + assertAssocInvariants av + assertAssocCompact av + assocSize av === length end + let endAv = mkAssoc end + assocFwd av === assocFwd endAv + assocBwd av === assocBwd endAv + +testAssocCases :: TestTree +testAssocCases = testGroup "Assoc case" (fmap testAssocCase allAssocCases) + +testAssocUnit :: TestTree +testAssocUnit = testUnit "Assoc unit" $ do + let a0 = assocNew :: AV + assertAssocInvariants a0 + assertAssocCompact a0 + assocSize a0 === 0 + let aKey = ENodeId 0 + aVal = toV 'a' + bKey = ENodeId 1 + bVal = toV 'b' + cKey = ENodeId 2 + cVal = toV 'c' + let members = [(aKey, aVal), (bKey, bVal), (cKey, cVal)] + let a1 = execState (for_ members (uncurry assocInsert)) a0 + assertAssocInvariants a1 + assertAssocCompact a0 + assocSize a1 === 3 + assocRoots a1 === [aKey, bKey, cKey] + assocLeaves a1 === [] + let (res, a2) = runState (assocInsert aKey bVal) a1 + res === (aKey, AssocInsertResMerged (ILS.singleton bKey)) + assertAssocInvariants a2 + assert $ assocCanCompact a2 + assocSize a2 === 2 + assocRoots a2 === [aKey, cKey] + assocLeaves a2 === [bKey] + let a3 = execState assocCompact a2 + assertAssocInvariants a3 + assertAssocCompact a3 + assocSize a3 === 2 + assocRoots a3 === [aKey, cKey] + assocLeaves a3 === [] + +type EGA = EGraph () ArithF +type EGP = Pat ArithF String + +testEgUnit :: TestTree +testEgUnit = after AllSucceed "Assoc unit" $ testUnit "EG unit" $ runS egNew $ do + -- We're going to have our egraph track the equality `2 + 2 = 4`. + -- We disable analysis + let ana = noAnalysis + -- Some simple terms: + let termFour = ArithConst 4 + termTwo = ArithConst 2 + termPlus = ArithPlus termTwo termTwo + -- And a simple pattern: + let pat = FreeEmbed (ArithPlusF (FreePure "x") (FreePure "y")) :: EGP + -- Test that the empty egraph is sane + testS $ \eg -> do + egClassSize eg === 0 + egNodeSize eg === 0 + -- Nothing is matched + testS $ \eg -> + match pat eg === [] + -- Add the term `4` + cidFour <- applyTestS (egAddTerm ana termFour) $ \(c, x) eg -> do + c === ChangedYes + egFindTerm termFour eg === Just x + egClassSize eg === 1 + egNodeSize eg === 1 + pure x + -- Add the term `2` + cidTwo <- applyTestS (egAddTerm ana termTwo) $ \(c, x) eg -> do + c === ChangedYes + x /== cidFour + egFindTerm termTwo eg === Just x + egClassSize eg === 2 + egNodeSize eg === 2 + pure x + -- Add the term `4` again and assert things haven't changed + applyTestS (egAddTerm ana termFour) $ \(c, x) eg -> do + c === ChangedNo + x === cidFour + egFindTerm termFour eg === Just x + egClassSize eg === 2 + egNodeSize eg === 2 + -- Still, nothing is matched + testS $ \eg -> + match pat eg === [] + -- Add the term `2 + 2` + cidPlus <- applyTestS (egAddTerm ana termPlus) $ \(c, x) eg -> do + c === ChangedYes + x /== cidFour + x /== cidTwo + egFindTerm termPlus eg === Just x + egClassSize eg === 3 + egNodeSize eg === 3 + pure x + -- We now match `2 + 2` + testS $ \eg -> + match pat eg === + [ MatchSubst + (Match cidPlus + (MatchPatEmbed + (ArithPlusF + (Match cidTwo (MatchPatPure "x")) + (Match cidTwo (MatchPatPure "y")) + ) + ) + ) + (HashMap.fromList [("x", cidTwo), ("y", cidTwo)]) + ] + -- Merge `4` and `4` and assert things haven't changed + applyTestS (egMerge cidFour cidFour) $ \m _ -> do + case m of + MergeResultUnchanged -> pure () + _ -> fail "expected unchanged merge" + -- Merge `2 + 2` and `4` + applyTestS (egMerge cidPlus cidFour) $ \m eg -> do + case m of + MergeResultChanged _ -> pure () + _ -> fail "expected changed merge" + egFindTerm termFour eg === Just cidFour + egFindTerm termPlus eg === Just cidFour + egFindTerm termTwo eg === Just cidTwo + -- We still match `2 + 2`, but the class is different + testS $ \eg -> + match pat eg === + [ MatchSubst + (Match cidFour + (MatchPatEmbed + (ArithPlusF + (Match cidTwo (MatchPatPure "x")) + (Match cidTwo (MatchPatPure "y")) + ) + ) + ) + (HashMap.fromList [("x", cidTwo), ("y", cidTwo)]) + ] + +type EGD = Max V +type EGF = BinTreeF V +type EGT = BinTree V +type EGV = EGraph EGD EGF + +maxVAnalysis :: EAnalysis EGD EGF +maxVAnalysis = \case + BinTreeLeafF v -> Max v + BinTreeBranchF d1 d2 -> d1 <> d2 + +assertEgInvariants :: (MonadTest m, Traversable f, Eq (f EClassId), Hashable (f EClassId), Show (f EClassId)) => EGraph d f -> m () +assertEgInvariants eg = do + let assoc = egNodeAssoc eg + hc = egHashCons eg + bwd = assocBwd assoc + rootNodes = ILS.fromList (assocRoots assoc) + leafNodes = ILS.fromList (assocLeaves assoc) + allNodes = ILS.union rootNodes leafNodes + ef = egEquivFind eg + rootClasses = ILS.fromList (efRoots ef) + leafClasses = ILS.fromList (efLeaves ef) + cm = egClassMap eg + cmClasses = ILS.fromList (ILM.keys cm) + -- Assert that root nodes and leaf nodes are disjoint + ILS.intersection rootNodes leafNodes === ILS.empty + -- Assert that root classes and leaf classes are disjoint + ILS.intersection rootClasses leafClasses === ILS.empty + -- Assert that the assoc is 1-1 etc + assertAssocInvariants assoc + -- Assert that the hashcons and assoc have equal key sets + ILS.fromList (ILM.keys hc) === allNodes + -- Assert that hashcons has exactly the same values as unionfind roots for all nodes + for_ (ILM.elems hc) $ \c -> + assert $ ILS.member c rootClasses + -- Assert that classmap only contains unionfind roots + cmClasses === rootClasses + -- For every node, assert in the nodes of some class + for_ (ILM.toList hc) $ \(n, c) -> do + let nodes = eciNodes (ILM.partialLookup c cm) + assert (assocMember n nodes) + -- For every root, assert is in all parent classes + for_ (ILM.toList hc) $ \(n, c) -> + when (ILS.member n rootNodes) $ do + -- for all children that are not of the node's own class + let children = ILS.filter (/= c) (foldMap ILS.singleton (assocPartialLookupByKey n assoc)) + for_ (ILS.toList children) $ \y -> do + -- look up child and assert in child's parents + let parents = eciParents (ILM.partialLookup y cm) + assert (ILS.member n parents) + -- For every class + cmNodes <- flipFoldM ILS.empty (ILM.toList cm) $ \accNodesSet (c, eci) -> do + let nodes = eciNodes eci + nodesSet = ILS.fromList (assocMembers nodes) + parents = eciParents eci + -- Assert that classmap node values are non-empty + nodesSet /== ILS.empty + -- Assert that classmap class has node values that are hashconsed to class + for_ (ILS.toList nodesSet) $ \n -> do + ILM.lookup n hc === Just c + -- Assert that classmap class has NO parents that are hashconsed to class + for_ (ILS.toList parents) $ \p -> + ILM.lookup p hc /== Just c + -- Assert we haven't seen these nodes before + assert $ ILS.disjoint nodesSet accNodesSet + -- Assert that the nodes and parents are disjoint + assert $ ILS.disjoint nodesSet parents + pure (ILS.union accNodesSet nodesSet) + let hcNodes = ILS.fromList (ILM.keys hc) + -- Assert hc keys are exactly the class nodes + cmNodes === hcNodes + -- Now test recanonicalization - we already know assoc fwd and bwd are 1-1 + for_ (HashMap.toList bwd) $ \(fc, _) -> + let recanon = evalState (egCanonicalize fc) eg + in recanon === Right fc + +data EgRound = EgRound + { egRoundTerms :: ![EGT] + , egRoundSets :: ![[EGT]] + , egRoundEqTests :: ![[EGT]] + , egRoundNeqTests :: ![(EGT, EGT)] + } deriving stock (Eq, Show) + +data EgCase = EgCase + { egCaseName :: !String + , egCaseRounds :: ![EgRound] + } deriving stock (Eq, Show) + +allEgCases :: [EgCase] +allEgCases = + let leafA = BinTreeLeaf (toV 'a') + leafB = BinTreeLeaf (toV 'b') + leafC = BinTreeLeaf (toV 'c') + leafD = BinTreeLeaf (toV 'd') + leafE = BinTreeLeaf (toV 'e') + leafTerms = [leafA, leafB, leafC, leafD] + parentAA = BinTreeBranch leafA leafA + parentAB = BinTreeBranch leafA leafB + parentAC = BinTreeBranch leafA leafC + parentAD = BinTreeBranch leafA leafD + parentBD = BinTreeBranch leafB leafD + parentCA = BinTreeBranch leafC leafA + simpleParentTerms = [parentAC, parentAD] + complexParentTerms = [parentAC, parentBD] + grandparentAAC = BinTreeBranch leafA parentAC + grandparentAAD = BinTreeBranch leafA parentAD + grandparentBAC = BinTreeBranch leafB parentAC + grandparentEAD = BinTreeBranch leafE parentAD + simpleGrandparentTerms = [grandparentAAC, grandparentAAD] + complexGrandparentTerms = [grandparentBAC, grandparentEAD] + in [ EgCase "simple" + [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC)] + , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafB, leafC)] + ] + , EgCase "transitive one round" + [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)] + , EgRound [] [[leafA, leafB], [leafB, leafC]] [[leafA, leafB, leafC]] [(leafA, leafD)] + ] + , EgCase "transitive two round" + [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)] + , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafA, leafD)] + , EgRound [] [[leafB, leafC]] [[leafA, leafB, leafC]] [(leafA, leafD)] + ] + , EgCase "simple parents" + [ EgRound simpleParentTerms [] [] [(leafC, leafD), (parentAC, parentAD)] + , EgRound [] [[leafC, leafD]] [[parentAC, parentAD]] [] + ] + , EgCase "complex parents one round" + [ EgRound complexParentTerms [] [] [(leafA, leafB), (leafC, leafD), (parentAC, parentBD)] + , EgRound [] [[leafA, leafB], [leafC, leafD]] [[leafA, leafB], [leafC, leafD], [parentAC, parentBD]] [] + ] + , EgCase "complex parents two round" + [ EgRound complexParentTerms [] [] [(leafA, leafB), (leafC, leafD), (parentAC, parentBD)] + , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafC, leafD), (parentAC, parentBD)] + , EgRound [] [[leafC, leafD]] [[leafA, leafB], [leafC, leafD], [parentAC, parentBD]] [] + ] + , EgCase "simple grandparents" + [ EgRound simpleGrandparentTerms [] [] [(leafC, leafD), (parentAC, parentAD), (grandparentAAC, grandparentAAD)] + , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD], [grandparentAAC, grandparentAAD]] [] + ] + , EgCase "complex grandparents bottom up" + [ EgRound complexGrandparentTerms [] [] [(leafC, leafD), (leafB, leafE), (parentAC, parentAD), (grandparentBAC, grandparentEAD)] + , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD]] [(leafB, leafE), (grandparentBAC, grandparentEAD)] + , EgRound [] [[leafB, leafE]] [[leafC, leafD], [leafB, leafE], [parentAC, parentAD], [grandparentBAC, grandparentEAD]] [] + ] + , EgCase "complex grandparents top down" + [ EgRound complexGrandparentTerms [] [] [(leafC, leafD), (leafB, leafE), (parentAC, parentAD), (grandparentBAC, grandparentEAD)] + , EgRound [] [[leafB, leafE]] [[leafB, leafE]] [(leafC, leafD), (parentAC, parentAD), (grandparentBAC, grandparentEAD)] + , EgRound [] [[leafC, leafD]] [[leafC, leafD], [leafB, leafE], [parentAC, parentAD], [grandparentBAC, grandparentEAD]] [] + ] + , EgCase "connect" + [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)] + , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafA, leafD)] + , EgRound [] [[leafC, leafD]] [[leafA, leafB], [leafC, leafD]] [(leafA, leafD)] + , EgRound [] [[leafB, leafD]] [[leafA, leafB, leafC, leafD]] [] + ] + , EgCase "mid grandparents" + [ EgRound simpleGrandparentTerms [] [] [(leafC, leafD), (parentAC, parentAD), (grandparentAAC, grandparentAAD)] + , EgRound [] [[parentAC, parentAD]] [[parentAC, parentAD], [grandparentAAC, grandparentAAD]] [(leafC, leafD)] + , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD], [grandparentAAC, grandparentAAD]] [] + ] + , EgCase "unify node" + [ EgRound [BinTreeBranch parentAC leafA, parentAA] [] [] [(parentAC, parentAA)] + , EgRound [] [[leafA, leafC]] [[parentAC, parentAA]] [] + ] + , EgCase "self parent" + [ EgRound [BinTreeBranch parentAC leafB] [] [] [] + , EgRound [] [[parentAC, leafA]] [] [] + ] + , EgCase "self parent again" + [ EgRound [leafB, parentAA] [] [] [] + , EgRound [] [[leafB, leafA], [leafB, parentAA]] [] [] + ] + , EgCase "dead add parent" + [ EgRound [parentAC, leafB] [[leafA, parentAC], [leafA, leafC]] [[parentAC, leafC]] [(leafA, leafB)] + , EgRound [parentCA] [] [[parentCA, parentAC]] [] + ] + , EgCase "repro 1" + [ EgRound [BinTreeBranch parentAA parentAB] [[leafB, leafA]] [] [(leafB, parentAA)] + , EgRound [leafA] [[parentAA, leafA]] [[leafB, parentAA]] [] + ] + , let grandparent = BinTreeBranch parentAA leafB + greatGrandparent = BinTreeBranch grandparent leafA + in EgCase "repro 2" + [ EgRound [greatGrandparent, leafA, leafC] [[greatGrandparent, parentAA], [leafA, leafB]] [] [] + , EgRound [leafA, leafA] [[leafA, grandparent]] [] [] + , EgRound [leafA, leafA, leafA] [[parentAA, leafA]] [] [] + ] + ] + +testEgCase :: EgCase -> TestTree +testEgCase (EgCase name rounds) = kase where + findMayTerm t = fmap (egFindTerm t) get + findTerm t = fmap fromJust (findMayTerm t) + findTerms ts = fmap ILS.fromList (for ts findTerm) + assertTermFound t = findMayTerm t >>= \mi -> assert (isJust mi) + assertTermsFound ts = for_ ts assertTermFound + kase = testUnit name $ runS egNew $ do + -- for each round + for_ rounds $ \(EgRound start act endEq endNeq) -> do + -- add initial terms and assert invariants hold + applyS (for_ start (egAddTerm maxVAnalysis)) + -- liftIO (putStrLn "===== post add =====") + -- testS $ liftIO . pPrint + testS assertEgInvariants + -- assert that all mentioned terms are in the egraph + assertTermsFound start + for_ act assertTermsFound + for_ endEq assertTermsFound + for_ endNeq $ \(x, y) -> do + -- also assert that neq terms are not themselves equal + x /== y + assertTermFound x + assertTermFound y + -- merge sets of terms and rebuild + applyS $ do + sets <- for act findTerms + mr <- egMergeMany (Seq.fromList sets) + case mr of + MergeResultMissing _ -> error "bad set" + _ -> pure () + -- assert invariants hold + testS assertEgInvariants + -- find merged terms again and assert they are in same classes + sets <- applyS $ for act findTerms + for_ sets $ \set -> ILS.size set === 1 + -- find final eq terms and assert they are in same classes + for_ endEq $ \ts -> do + set <- applyS (findTerms ts) + ILS.size set === 1 + -- find final neq terms and assert they are not in same class + for_ endNeq $ \(x, y) -> do + i <- applyS (findTerm x) + j <- applyS (findTerm y) + i /== j + +testEgCases :: TestTree +testEgCases = testGroup "Eg case" $ + testEgCase <$> allEgCases + +testEgNew :: TestTree +testEgNew = testUnit "EG new" $ do + eg0 <- fullyEvaluate (egNew :: EGV) + egNodeSize eg0 === 0 + egClassSize eg0 === 0 + assertEgInvariants eg0 + +genNodePairs :: Range Int -> EGV -> Gen [(EClassId, EClassId)] +genNodePairs nOpsRange eg = genListOfDistinctPairs nOpsRange (ILM.keys (egClassMap eg)) + +genSomeList :: [a] -> Gen [a] +genSomeList xs = go where + go = Gen.recursive Gen.choice [Gen.constant [], fmap pure (Gen.element xs)] [Gen.subterm2 go go (++)] + +genBinTree :: Gen a -> Gen (BinTree a) +genBinTree genA = genEither where + genLeaf = fmap BinTreeLeaf genA + genBranch = Gen.subterm2 genEither genEither BinTreeBranch + genEither = Gen.recursive Gen.choice [genLeaf] [genBranch] + +genBinTreeMembers :: Int -> Gen [BinTree V] +genBinTreeMembers maxElems = Gen.list (Range.linear 0 maxElems) (genBinTree (genV maxElems)) + +-- An alternative to 'genBinTreeMembers' that makes smaller trees +mkSimpleTreeLevels :: Int -> [BinTree V] +mkSimpleTreeLevels maxElems = + let letters = take maxElems (['a'..'z'] ++ ['A'..'Z']) + zeroLevel = fmap (BinTreeLeaf . toV) letters + mkLevel y x = (BinTreeBranch <$> x <*> y) ++ (BinTreeBranch <$> y <*> x) + mkLevels y = foldr (\x r -> mkLevel y x ++ r) (BinTreeBranch <$> y <*> y) + oneLevel = mkLevels zeroLevel [] + twoLevel = mkLevels oneLevel [zeroLevel] + anyLevel = zeroLevel ++ oneLevel ++ twoLevel + in anyLevel + +testEgProp :: TestLimit -> TestTree +testEgProp lim = after AllSucceed "EG unit" $ after AllSucceed "EG cases" $ testProp "EG prop" lim prop where + maxElems = 10 + termGen = genBinTreeMembers maxElems + -- Guarantee yourself small trees with this: + -- termGen = genSomeList (mkSimpleTreeLevels maxElems) + prop = do + eg0 <- fullyEvaluate (egNew :: EGV) + rounds <- forAll (Gen.element [1, 2, 3]) + body rounds eg0 + body (rounds :: Int) eg0 = do + members <- forAll termGen + let nMembers = length members + nOpsRange = Range.linear 0 (nMembers * nMembers) + eg1 <- fullyEvaluate (execState (for_ members (egAddTerm maxVAnalysis)) eg0) + assertEgInvariants eg1 + pairs <- forAll (genNodePairs nOpsRange eg1) + let merge = do + mr <- egMergeMany (Seq.fromList (fmap (\(a, b) -> ILS.fromList [a, b]) pairs)) + case mr of + MergeResultMissing _ -> error "bad set" + _ -> pure () + eg2 <- fullyEvaluate (execState merge eg1) + assertEgInvariants eg2 + unless (rounds == 1) (body (rounds - 1) eg2) + +type M = IntLikeMap ENodeId Char + +testILM :: TestTree +testILM = testUnit "ILM unit" $ do + let mLeft = ILM.fromList [(ENodeId 0, 'a'), (ENodeId 1, 'b')] :: M + mRight = ILM.fromList [(ENodeId 1, 'x'), (ENodeId 2, 'c')] :: M + mMerged = ILM.fromList [(ENodeId 0, 'a'), (ENodeId 1, 'b'), (ENodeId 2, 'c')] :: M + mLeft <> mRight === mMerged + +main :: IO () +main = testMain $ \lim -> testGroup "Overeasy" + [ testILM + , testEfUnit + , testAssocUnit + , testAssocCases + , testEgUnit + , testEgNew + , testEgCases + , testEfProp lim + , testEgProp lim + ] diff --git a/test/Test/Overeasy/Arith.hs b/test/Test/Overeasy/Arith.hs new file mode 100644 --- /dev/null +++ b/test/Test/Overeasy/Arith.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE TemplateHaskell #-} + +module Test.Overeasy.Arith + ( ArithF (..) + , Arith (..) + ) where + +import Control.DeepSeq (NFData) +import Data.Functor.Foldable.TH (makeBaseFunctor) +import Data.Hashable (Hashable) +import GHC.Generics (Generic) + +data Arith = + ArithPlus Arith Arith + | ArithTimes Arith Arith + | ArithShiftL Arith !Int + | ArithShiftR Arith !Int + | ArithConst !Int + deriving stock (Eq, Show, Generic) + deriving anyclass (Hashable, NFData) + +-- Generates 'ArithF' and other recursion-schemes boilerplate +makeBaseFunctor ''Arith + +deriving stock instance Eq a => Eq (ArithF a) +deriving stock instance Show a => Show (ArithF a) +deriving stock instance Generic (ArithF a) +deriving anyclass instance Hashable a => Hashable (ArithF a) +deriving anyclass instance NFData a => NFData (ArithF a) diff --git a/test/Test/Overeasy/BinTree.hs b/test/Test/Overeasy/BinTree.hs new file mode 100644 --- /dev/null +++ b/test/Test/Overeasy/BinTree.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE DeriveAnyClass #-} + +module Test.Overeasy.BinTree + ( BinTreeF (..) + , BinTree (..) + , pattern BinTreeLeaf + , pattern BinTreeBranch + ) where + +import Control.Applicative (liftA2) +import Control.DeepSeq (NFData) +import Data.Bifoldable (Bifoldable (..)) +import Data.Bifunctor (Bifunctor (..)) +import Data.Bitraversable (Bitraversable (..)) +import Data.Functor.Foldable (Base, Corecursive (..), Recursive (..)) +import Data.Hashable (Hashable) +import GHC.Generics (Generic) + +data BinTreeF a r = + BinTreeLeafF !a + | BinTreeBranchF r r + deriving stock (Eq, Ord, Show, Functor, Foldable, Traversable, Generic) + deriving anyclass (Hashable, NFData) + +instance Bifunctor BinTreeF where + bimap f g = \case + BinTreeLeafF a -> BinTreeLeafF (f a) + BinTreeBranchF x y -> BinTreeBranchF (g x) (g y) + +instance Bifoldable BinTreeF where + bifoldr f g z = \case + BinTreeLeafF a -> f a z + BinTreeBranchF x y -> g x (g y z) + +instance Bitraversable BinTreeF where + bitraverse f g = \case + BinTreeLeafF a -> fmap BinTreeLeafF (f a) + BinTreeBranchF x y -> liftA2 BinTreeBranchF (g x) (g y) + +newtype BinTree a = BinTree { unBinTree :: BinTreeF a (BinTree a) } + deriving newtype (Eq, Ord, Show, Hashable, NFData) + +pattern BinTreeLeaf :: a -> BinTree a +pattern BinTreeLeaf a = BinTree (BinTreeLeafF a) + +pattern BinTreeBranch :: BinTree a -> BinTree a -> BinTree a +pattern BinTreeBranch a b = BinTree (BinTreeBranchF a b) + +{-# COMPLETE BinTreeLeaf, BinTreeBranch #-} + +instance Functor BinTree where + fmap f = go where + go = BinTree . bimap f go . unBinTree + +instance Foldable BinTree where + foldr f z0 x0 = go x0 z0 where + go x z = bifoldr f go z (unBinTree x) + +instance Traversable BinTree where + traverse f = go where + go = fmap BinTree . bitraverse f go . unBinTree + +type instance Base (BinTree a) = BinTreeF a + +instance Recursive (BinTree a) where + project = unBinTree + +instance Corecursive (BinTree a) where + embed = BinTree