-- Copyright 2022 Google LLC -- -- Use of this source code is governed by a BSD-style -- license that can be found in the LICENSE file or at -- https://developers.google.com/open-source/licenses/bsd module OccAnalysisSpec (spec) where import Prelude hiding (unlines) import Data.Maybe (catMaybes) import Data.Text import Test.Hspec import ConcreteSyntax (parseUModule) import AbstractSyntax (parseExpr) import Err import Inference (inferTopUExpr, synthTopE) import Name import OccAnalysis import Occurrence import Simplify import SourceRename (renameSourceNamesUExpr) import Types.Core import Types.Imp (Backend (..)) import Types.Primitives import Types.Source import TopLevel import QueryType sourceTextToBlocks :: (Topper m, Mut n) => Text -> m n [SBlock n] sourceTextToBlocks source = do let (UModule _ deps sourceBlocks) = parseUModule Main source mapM_ ensureModuleLoaded deps catMaybes <$> mapM sourceBlockToBlock sourceBlocks sourceBlockToBlock :: (Topper m, Mut n) => SourceBlock -> m n (Maybe (SBlock n)) sourceBlockToBlock block = case sbContents block of Misc (ImportModule moduleName) -> importModule moduleName >> return Nothing Command (EvalExpr (Printed _)) expr -> Just <$> (parseExpr expr >>= uExprToBlock) UnParseable _ s -> throw ParseErr s _ -> error $ "Unexpected SourceBlock " ++ pprint block ++ " in unit tests" uExprToBlock :: (Topper m, Mut n) => UExpr 'VoidS -> m n (SBlock n) uExprToBlock expr = do renamed <- renameSourceNamesUExpr expr typed <- inferTopUExpr renamed synthed <- synthTopE typed SimplifiedTopLam (TopLam _ _ (LamExpr Empty block)) (CoerceRecon _) <- simplifyTopBlock synthed return block findRunIOAnnotation :: SBlock n -> LetAnn findRunIOAnnotation (Abs decls _) = go decls where go :: Nest SDecl n l -> LetAnn go (Nest (Let _ (DeclBinding ann (PrimOp (Hof (TypedHof _ (RunIO _)))))) _) = ann go (Nest _ rest) = go rest go Empty = error "RunIO not found" analyze :: EvalConfig -> TopStateEx -> [Text] -> IO LetAnn analyze cfg env code = fst <$> runTopperM cfg env do [block] <- sourceTextToBlocks $ unlines code lam <- asTopLam $ LamExpr Empty block TopLam _ _ (LamExpr Empty block') <- analyzeOccurrences lam -- The RunIO is generated by simplifying `unreachable()` in the examples -- below. If we want compound examples that have more than one RunIO block, -- we will need better pattern-matching. return $ findRunIOAnnotation block'; spec :: Spec spec = do let cfg = EvalConfig LLVM [LibBuiltinPath] Nothing Nothing Nothing Optimize PrintCodegen -- or just initTopState, to always compile the prelude during unit tests? init_env <- runIO loadCache (_, env) <- runIO $ runTopperM cfg init_env $ ensureModuleLoaded Prelude describe "Occurrence analysis" do it "counts a reference as a use" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " xs" ] ann `shouldBe` OccInfoPure (UsageInfo One (0, One)) it "counts indexing in a for as one use" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " for i. xs[i]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) it "counts indexing depth in nested fors" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => (Fin 3) => Float = unreachable()" , " for i j. xs[i, j]" ] ann `shouldBe` OccInfoPure (UsageInfo One (2, One)) it "counts two array uses" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " (for i. xs[i], for j. xs[j])" ] ann `shouldBe` OccInfoPure (UsageInfo Two (1, Two)) it "counts array and non-array uses" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " (for i. xs[i], xs)" ] ann `shouldBe` OccInfoPure (UsageInfo Two (1, Two)) it "counts different case arms as static but not dynamic uses" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " if unreachable()" , " then for i. xs[i]" , " else for j. xs[j]" ] ann `shouldBe` OccInfoPure (UsageInfo Two (1, One)) it "understands one index injection" do ann <- analyze cfg env [ ":p" , " xs : Either(Fin 10, Fin 4) => Float = unreachable()" , " for j. xs[Left j]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) it "understands distinct index injections" do ann <- analyze cfg env [ ":p" , " xs : Either(Fin 10, Fin 4) => Float = unreachable()" , " (for i. xs[Left i], for j. xs[Right j])" ] ann `shouldBe` OccInfoPure (UsageInfo Two (1, One)) it "detects and eschews index arithmetic" do ann <- analyze cfg env [ ":p" , " xs : (Fin 4) => Float = unreachable()" , " for i:(Fin 3). xs[(ordinal i + 1)@_]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "detects non-nested single-uses cases despite index arithmetic" do ann <- analyze cfg env [ ":p" , " xs : (Fin 4) => Float = unreachable()" , " xs[1@_]" ] -- Arguably, should be able to prove that zero levels of exposed indexing -- (not one) suffice for inlining xs to be safe here, but the current -- occurrence analysis doesn't prove it yet. ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) it "detects nested single-uses cases despite index arithmetic" do ann <- analyze cfg env [ ":p" , " xs : (Fin 4) => (Fin 3) => Float = unreachable()" , " for i:(Fin 3). xs[(ordinal i + 1)@_, i]" ] ann `shouldBe` OccInfoPure (UsageInfo One (2, One)) it "detects repeated access" do ann <- analyze cfg env [ ":p" , " xs : (Fin 4) => Float = unreachable()" , " for i j:(Fin 5). xs[i]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "does not count the `trace` pattern as repeated access" do ann <- analyze cfg env [ ":p" , " xs : (Fin 4) => (Fin 4) => Float = unreachable()" , " for i. xs[i, i]" ] -- Arguably, should be able to prove that only one level of exposed -- indexing (not two) suffice for inlining xs to be safe here, but doesn't -- prove it yet. ann `shouldBe` OccInfoPure (UsageInfo One (2, One)) it "solves safe sum-over-max" do ann <- analyze cfg env [ ":p" , " xs : Either(Either(Fin 4, Fin 4), Either(Fin 4, Fin 4)) => Float = unreachable()" , " ys = for i." , " if unreachable()" , " then xs[Left (Left i)]" , " else xs[Left (Right i)]" , " zs = for j." , " if unreachable()" , " then xs[Right (Left j)]" , " else xs[Right (Right j)]" , " (ys, zs)" ] ann `shouldBe` OccInfoPure (UsageInfo (Bounded 4) (1, One)) it "solves unsafe sum-over-max" do ann <- analyze cfg env [ ":p" , " xs : Either(Either(Fin 4, Fin 4), Either(Fin 4, Fin 4)) => Float = unreachable()" , " ys = for i." , " if unreachable()" , " then xs[Left (Left i)]" , " else xs[Left (Right i)]" , " zs = for j." , " if unreachable()" , " then xs[Right (Left j)]" , " else xs[Left (Left j)]" , " (ys, zs)" ] -- One of the code paths hits the same elements(s) ann `shouldBe` OccInfoPure (UsageInfo (Bounded 4) (1, Two)) it "does not penalize referring to indices in scope" do ann <- analyze cfg env [ ":p" , " j = 1@(Fin 3)" , " xs : (Fin 10) => (Fin 3) => Float = unreachable()" , " for i. xs[i, j]" ] -- Arguably, should be able to prove that only one level of exposed -- indexing (not two) suffice for inlining xs to be safe here, but doesn't -- prove it yet. ann `shouldBe` OccInfoPure (UsageInfo One (2, One)) it "is conservative about potential collisions between indices in scope" do ann <- analyze cfg env [ ":p" , " j = 1@(Fin 3)" , " k = 1@(Fin 3)" , " xs : (Fin 10) => (Fin 3) => Float = unreachable()" , " (for i. xs[i, j], for i. xs[i, k])" ] ann `shouldBe` OccInfoPure (UsageInfo Two (2, Two)) it "is not confused by potential collisions at an early indexing depth" do ann <- analyze cfg env [ ":p" , " j = 1@(Fin 10)" , " k = 1@(Fin 10)" , " xs : (Fin 10) => (Fin 3) => Float = unreachable()" , " (for i. xs[j, i], for i. xs[k, i])" ] ann `shouldBe` OccInfoPure (UsageInfo Two (2, Two)) it "does not crash on indexing by case-bound binders" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " for i." , " case i of" , " (Left j) -> xs[j]" , " (Right k) -> xs[k]" ] -- TODO actually, it should be possible to get this to be Bounded 2 rather -- than Unbounded. We get Unbounded because `occAlt` assumes that the -- binder of a case (in this case `j`) is an "unknown function" of the -- scrutinee. That's conservative, but in reality the function is very -- well known, and even injective, but just not total (and not the same -- across arms of the `case`). However, trying to fix that would unmask -- another bug, which is that mapping `case` to `max` is only correct if -- the scrutinee doesn't depend on any binders being iterated. In fact, -- in this example it does, so both arms of the `case` end up being taken, -- albeit at different iterations of the `i` loop. To analyze this -- correctly, we would need to know that `j` and `k` may collide across -- `case` arms, though not within an arm. ann `shouldBe` OccInfoPure (UsageInfo Two (1, Unbounded)) it "does not crash on indexing by state-effect-bound binders" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " with_state (0 @ Fin 10) \\ref." , " xs[get ref]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) it "assumes state references touch everything" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " with_state (0 @ Fin 10) \\ref." , " for i:(Fin 3)." , " xs[get ref]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "assumes state references touch everything even if the initializer doesn't" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " for i:(Fin 10)." , " with_state i \\ref." , " xs[get ref]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "analyzes through accum effects" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " run_accum (AddMonoid Float) \\ref." , " for i." , " ref += xs[i]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) -- TODO Should probably construct an example of indexing with the ref of a -- run_accum (and, for that matter, with the handle of a run_state or -- run_accum), but I'm not sure how to make either of those well-typed. it "does not crash on indexing by reader-effect-bound binders" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" , " with_reader (0 @ Fin 10) \\ref." , " xs[ask ref]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, One)) it "understands that while loops access repeatedly" do ann <- analyze cfg env [ ":p" , " xs : (Fin 10) => Float = unreachable()" -- with_state prevents the access from appearing dead , " with_state 0 \\ref." , " while \\." , " for i." , " ref := (get ref) + xs[i]" , " False" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "does not crash on index-defining bindings" do ann <- analyze cfg env [ ":p" , " xs : Either(Fin 10, Fin 3) => Float = unreachable()" , " for i:(Fin 10)." , " j = Left(unsafe_from_ordinal(1 + ordinal i))" , " xs[j]" ] ann `shouldBe` OccInfoPure (UsageInfo One (1, Unbounded)) it "understands indexing by literals" do ann <- analyze cfg env [ ":p" , " xs : Either(Fin 10, Fin 3) => Float = unreachable()" , " (xs[0@_], xs[0@_])" ] ann `shouldBe` OccInfoPure (UsageInfo Two (1, Two)) it "is conservative about distict literal indices" do ann <- analyze cfg env [ ":p" , " xs : Either(Fin 10, Fin 3) => Float = unreachable()" , " (xs[0@_], xs[1@_])" ] -- TODO In this case, we should be able to detect non-collision of -- indexing by 0 and by 1; but assuming they may collide is safe. ann `shouldBe` OccInfoPure (UsageInfo Two (1, Two)) it "is conservative about ix dicts" do ann <- analyze cfg env [ ":p" , " n : Nat = unreachable()" , " xs : (Fin n) => Nat = iota (Fin n)" , " sum xs" ] -- We consider `n` statically unbounded because we assume that its ix dict -- may be inlined uncontrollably later. ann `shouldBe` OccInfoPure (UsageInfo Unbounded (0, Unbounded))