Differentiate between Import and Superclass edge tests

src/Data/Abstract/ScopeGraph.hs

@@ -196,10 +196,10 @@ lookupDeclaration declaration scope g = do
   index <- Seq.findIndexR (((Declaration declaration) ==) . fst) dataSeq
   (, Position index) <$> Seq.lookup index dataSeq
 
-declarationNames :: Ord address => Scope address -> ScopeGraph address -> [Declaration]
-declarationNames scope scopeGraph = localDeclarations <> edgeNames
-  where addresses = join (Map.elems $ Map.withoutKeys (edges scope) (Set.fromList [Lexical, Import, Export]))
-        edgeNames = addresses >>= toList . flip lookupScope scopeGraph >>= flip declarationNames scopeGraph
+declarationNames :: Ord address => [EdgeLabel] -> Scope address -> ScopeGraph address -> [Declaration]
+declarationNames edgeLabels scope scopeGraph = localDeclarations <> edgeNames
+  where addresses = join (Map.elems $ Map.restrictKeys (edges scope) (Set.fromList edgeLabels))
+        edgeNames = addresses >>= toList . flip lookupScope scopeGraph >>= flip (declarationNames edgeLabels) scopeGraph
         localDeclarations = toList . fmap fst $ declarations scope
 
 

test/Analysis/Go/Spec.hs

@@ -15,7 +15,7 @@ spec config = parallel $ do
       case ModuleTable.lookup "main.go" <$> res of
         Right (Just (Module _ (scopeAndFrame, valueRef) :| [])) -> do
           () <$ SpecHelpers.lookupDeclaration "foo" scopeAndFrame heap scopeGraph `shouldBe` Just ()
-          SpecHelpers.lookupObjectMembers "foo" scopeAndFrame heap scopeGraph `shouldBe` Just ["New"]
+          SpecHelpers.lookupMembers "foo" Import scopeAndFrame heap scopeGraph `shouldBe` Just ["New"]
           () <$ SpecHelpers.lookupDeclaration "main" scopeAndFrame heap scopeGraph `shouldBe` Just ()
           () <$ SpecHelpers.lookupDeclaration "Bar" scopeAndFrame heap scopeGraph `shouldBe` Just ()
           () <$ SpecHelpers.lookupDeclaration "Rab" scopeAndFrame heap scopeGraph `shouldBe` Just ()

test/Analysis/Python/Spec.hs

@@ -19,8 +19,8 @@ spec config = parallel $ do
           const () <$> SpecHelpers.lookupDeclaration "a" scopeAndFrame heap scopeGraph `shouldBe` Just ()
           const () <$> SpecHelpers.lookupDeclaration "b" scopeAndFrame heap scopeGraph `shouldBe` Just ()
 
-          fromJust (SpecHelpers.lookupObjectMembers "a" scopeAndFrame heap scopeGraph) `shouldContain` [ "foo" ]
-          fromJust (SpecHelpers.lookupObjectMembers "b" scopeAndFrame heap scopeGraph) `shouldContain` ["c"]
+          fromJust (SpecHelpers.lookupMembers "a" Import scopeAndFrame heap scopeGraph) `shouldContain` [ "foo" ]
+          fromJust (SpecHelpers.lookupMembers "b" Import scopeAndFrame heap scopeGraph) `shouldContain` ["c"]
           -- (derefQName heap ("b" :| ["c"]) env >>= deNamespace heap) `shouldBe` Just ("c", ["baz"])
         other -> expectationFailure (show other)
 
@@ -57,7 +57,7 @@ spec config = parallel $ do
         Right (Just (Module _ (scopeAndFrame, valueRef) :| [])) -> do
           () <$ SpecHelpers.lookupDeclaration "Foo" scopeAndFrame heap scopeGraph `shouldBe` Just ()
           () <$ SpecHelpers.lookupDeclaration "Bar" scopeAndFrame heap scopeGraph `shouldBe` Just ()
-          SpecHelpers.lookupObjectMembers "Bar" scopeAndFrame heap scopeGraph `shouldBe` Just [ "dang" ]
+          SpecHelpers.lookupMembers "Bar" Superclass scopeAndFrame heap scopeGraph `shouldBe` Just [ "dang" ]
           valueRef `shouldBe` Rval (String "\"bar\"")
         other -> expectationFailure (show other)
 
@@ -65,7 +65,7 @@ spec config = parallel $ do
       (scopeGraph, (heap, res)) <- evaluate ["multiple_inheritance.py"]
       case ModuleTable.lookup "multiple_inheritance.py" <$> res of
         Right (Just (Module _ (scopeAndFrame, valueRef) :| [])) -> do
-          SpecHelpers.lookupObjectMembers "Baz" scopeAndFrame heap scopeGraph `shouldBe` Just [ "dang" ]
+          SpecHelpers.lookupMembers "Baz" Superclass scopeAndFrame heap scopeGraph `shouldBe` Just [ "dang" ]
           valueRef `shouldBe` Rval (String "\"bar!\"")
         other -> expectationFailure (show other)
 

test/Analysis/TypeScript/Spec.hs

@@ -40,8 +40,8 @@ spec config = parallel $ do
           -- Env.names env `shouldBe` [ "b", "z" ]
           () <$ SpecHelpers.lookupDeclaration "b" scopeAndFrame heap scopeGraph `shouldBe` Just ()
           () <$ SpecHelpers.lookupDeclaration "z" scopeAndFrame heap scopeGraph `shouldBe` Just ()
-          lookupObjectMembers "b" scopeAndFrame heap scopeGraph `shouldBe` Just  [ "foo", "baz" ]
-          lookupObjectMembers "z" scopeAndFrame heap scopeGraph `shouldBe` Just  [ "foo", "baz" ]
+          lookupMembers "b" Import scopeAndFrame heap scopeGraph `shouldBe` Just  [ "foo", "baz" ]
+          lookupMembers "z" Import scopeAndFrame heap scopeGraph `shouldBe` Just  [ "foo", "baz" ]
           -- (Heap.lookupDeclaration "b" heap  >>= deNamespace heap) `shouldBe` Just ("b", [ "baz", "foo" ])
           -- (Heap.lookupDeclaration "z" heap >>= deNamespace heap) `shouldBe` Just ("z", [ "baz", "foo" ])
           () <$ SpecHelpers.lookupDeclaration "baz" scopeAndFrame heap scopeGraph `shouldBe` Nothing

test/SpecHelpers.hs

@@ -15,10 +15,12 @@ module SpecHelpers
 , LogQueue
 , StatQueue
 , lookupDeclaration
-, lookupObjectMembers
+, lookupMembers
+, EdgeLabel(..)
 ) where
 
 import Control.Abstract hiding (lookupDeclaration)
+import Data.Abstract.ScopeGraph (EdgeLabel(..))
 import qualified Data.Abstract.ScopeGraph as ScopeGraph
 import qualified Data.Abstract.Heap as Heap
 import Control.Arrow ((&&&))
@@ -161,26 +163,28 @@ testEvaluating
 type Val term = Value term Precise
 
 
-objectMembers :: Heap Precise Precise (Value term Precise)
-            -> ScopeGraph Precise
-            -> Value term Precise
-            -> Maybe [Name]
-objectMembers heap scopeGraph (Object frame) = frameNames heap scopeGraph frame
-objectMembers heap scopeGraph (Class _ _ frame) = frameNames heap scopeGraph frame
-objectMembers _ _ _                          = Nothing
+members :: EdgeLabel
+        -> Heap Precise Precise (Value term Precise)
+        -> ScopeGraph Precise
+        -> Value term Precise
+        -> Maybe [Name]
+members edgeLabel heap scopeGraph (Object frame)    = frameNames [ edgeLabel ] heap scopeGraph frame
+members edgeLabel heap scopeGraph (Class _ _ frame) = frameNames [ edgeLabel ] heap scopeGraph frame
+members _ _ _ _                                     = Nothing
 
-frameNames :: Heap Precise Precise (Value term Precise)
+frameNames :: [ EdgeLabel ]
+           -> Heap Precise Precise (Value term Precise)
            -> ScopeGraph Precise
            -> Precise
            -> Maybe [ Name ]
-frameNames heap scopeGraph frame = do
+frameNames edge heap scopeGraph frame = do
   scopeAddress <- Heap.scopeLookup frame heap
   scope <- ScopeGraph.lookupScope scopeAddress scopeGraph
-  pure (unDeclaration <$> ScopeGraph.declarationNames scope scopeGraph)
+  pure (unDeclaration <$> ScopeGraph.declarationNames edge scope scopeGraph)
 
-lookupObjectMembers :: Name -> (Precise, Precise) -> Heap Precise Precise (Value term Precise) -> ScopeGraph Precise -> Maybe [ Name ]
-lookupObjectMembers name scopeAndFrame heap scopeGraph =
-  (lookupDeclaration name scopeAndFrame heap scopeGraph >>= objectMembers heap scopeGraph . Prelude.head)
+lookupMembers :: Name -> EdgeLabel -> (Precise, Precise) -> Heap Precise Precise (Value term Precise) -> ScopeGraph Precise -> Maybe [ Name ]
+lookupMembers name edgeLabel scopeAndFrame heap scopeGraph =
+  (lookupDeclaration name scopeAndFrame heap scopeGraph >>= members edgeLabel heap scopeGraph . Prelude.head)
 
 lookupDeclaration :: Name -> (Precise, Precise) -> Heap Precise Precise (Value term Precise) -> ScopeGraph Precise -> Maybe [ Value term Precise ]
 lookupDeclaration name (currentScope, currentFrame) heap scopeGraph = do