2222import java .util .HashSet ;
2323import java .util .Iterator ;
2424import java .util .Map ;
25+ import java .util .Optional ;
2526import java .util .Set ;
2627import java .util .stream .Collectors ;
28+ import java .util .stream .Stream ;
29+ import javax .annotation .CheckForNull ;
30+ import org .sonar .plugins .python .api .tree .Expression ;
31+ import org .sonar .plugins .python .api .tree .Name ;
32+ import org .sonar .plugins .python .api .tree .Tree ;
2733import org .sonar .python .semantic .v2 .SymbolV2 ;
34+ import org .sonar .python .semantic .v2 .SymbolV2Utils ;
2835import org .sonar .python .semantic .v2 .TypeTable ;
36+ import org .sonar .python .semantic .v2 .UsageV2 ;
37+ import org .sonar .python .tree .NameImpl ;
38+ import org .sonar .python .tree .TreeUtils ;
2939import org .sonar .python .types .v2 .PythonType ;
40+ import org .sonar .python .types .v2 .UnionType ;
3041
3142public class AstBasedPropagation {
3243 private final Map <SymbolV2 , Set <Propagation >> propagationsByLhs ;
33- private final TypeTable typeTable ;
44+ private final Propagator propagator ;
3445
3546 public AstBasedPropagation (Map <SymbolV2 , Set <Propagation >> propagationsByLhs , TypeTable typeTable ) {
3647 this .propagationsByLhs = propagationsByLhs ;
37- this .typeTable = typeTable ;
48+ this .propagator = new Propagator ( typeTable ) ;
3849 }
3950
4051 public Map <SymbolV2 , Set <PythonType >> processPropagations (Set <SymbolV2 > trackedVars ) {
@@ -56,18 +67,76 @@ public Map<SymbolV2, Set<PythonType>> processPropagations(Set<SymbolV2> trackedV
5667 return propagations .stream ().collect (Collectors .groupingBy (Propagation ::lhsSymbol , Collectors .mapping (Propagation ::rhsType , Collectors .toSet ())));
5768 }
5869
59- private static void applyPropagations (Set <Propagation > propagations , Set <SymbolV2 > initializedVars , boolean checkDependenciesReadiness ) {
70+ private void applyPropagations (Set <Propagation > propagations , Set <SymbolV2 > initializedVars , boolean checkDependenciesReadiness ) {
6071 Set <Propagation > workSet = new HashSet <>(propagations );
6172 while (!workSet .isEmpty ()) {
6273 Iterator <Propagation > iterator = workSet .iterator ();
6374 Propagation propagation = iterator .next ();
6475 iterator .remove ();
6576 if (!checkDependenciesReadiness || propagation .areDependenciesReady (initializedVars )) {
66- boolean learnt = propagation .propagate (initializedVars );
77+ boolean learnt = propagator .propagate (propagation , initializedVars );
6778 if (learnt ) {
6879 workSet .addAll (propagation .dependents ());
6980 }
7081 }
7182 }
7283 }
84+
85+ private record Propagator (TypeTable typeTable ) {
86+ /**
87+ * @return true if the propagation effectively changed the inferred type of assignment LHS
88+ */
89+ public boolean propagate (Propagation propagation , Set <SymbolV2 > initializedVars ) {
90+ PythonType rhsType = propagation .rhsType ();
91+ Name lhsName = propagation .lhsName ();
92+ SymbolV2 lhsSymbol = propagation .lhsSymbol ();
93+ if (initializedVars .add (lhsSymbol )) {
94+ propagateTypeToUsages (propagation , rhsType );
95+ return true ;
96+ } else {
97+ PythonType currentType = currentType (lhsName );
98+ if (currentType == null ) {
99+ return false ;
100+ }
101+ PythonType newType = UnionType .or (rhsType , currentType );
102+ propagateTypeToUsages (propagation , newType );
103+ return !newType .equals (currentType );
104+ }
105+ }
106+
107+ private void propagateTypeToUsages (Propagation propagation , PythonType newType ) {
108+ Tree scopeTree = propagation .scopeTree (propagation .lhsName ());
109+ getSymbolNonDeclarationUsageTrees (propagation .lhsSymbol )
110+ .filter (NameImpl .class ::isInstance )
111+ .map (NameImpl .class ::cast )
112+ // Avoid propagation to usages in nested scopes, as this may lead to FPs
113+ .filter (n -> isInSameScope (propagation , n , scopeTree ))
114+ .forEach (n -> n .typeV2 (newType ));
115+ }
116+
117+ @ CheckForNull
118+ private static PythonType currentType (Name lhsName ) {
119+ return Optional .ofNullable (lhsName .symbolV2 ())
120+ .stream ()
121+ .flatMap (Propagator ::getSymbolNonDeclarationUsageTrees )
122+ .flatMap (TreeUtils .toStreamInstanceOfMapper (Expression .class ))
123+ .findFirst ()
124+ .map (Expression ::typeV2 )
125+ .orElse (null );
126+ }
127+
128+ public static Stream <Tree > getSymbolNonDeclarationUsageTrees (SymbolV2 symbol ) {
129+ return symbol .usages ()
130+ .stream ()
131+ // Function and class definition names will always have FunctionType and ClassType respectively
132+ // so they are filtered out of type propagation
133+ .filter (u -> !SymbolV2Utils .isDeclaration (u ))
134+ .map (UsageV2 ::tree );
135+ }
136+
137+ private boolean isInSameScope (Propagation propagation , Name n , Tree scopeTree ) {
138+ return Optional .ofNullable (propagation .scopeTree (n )).filter (scopeTree ::equals ).isPresent ();
139+ }
140+ }
141+
73142}
0 commit comments