1 files changed, 12 insertions, 4 deletions
diff --git a/lib/Analysis/DataStructure/Steensgaard.cpp b/lib/Analysis/DataStructure/Steensgaard.cpp
index aa7f54b890..747ae2fc4b 100644
--- a/lib/Analysis/DataStructure/Steensgaard.cpp
+++ b/lib/Analysis/DataStructure/Steensgaard.cpp
@@ -167,14 +167,22 @@ bool Steens::runOnModule(Module &M) {
     }
   }
 
-  // Remove our knowledge of what the return values of the functions are.
-  ResultGraph->getReturnNodes().clear();
+  // Remove our knowledge of what the return values of the functions are, except
+  // for functions that are externally visible from this module (e.g. main).  We
+  // keep these functions so that their arguments are marked incomplete.
+  for (DSGraph::ReturnNodesTy::iterator I =
+         ResultGraph->getReturnNodes().begin(),
+         E = ResultGraph->getReturnNodes().end(); I != E; )
+    if (I->first->hasInternalLinkage())
+      ResultGraph->getReturnNodes().erase(I++);
+    else
+      ++I;
 
   // Update the "incomplete" markers on the nodes, ignoring unknownness due to
   // incoming arguments...
   ResultGraph->maskIncompleteMarkers();
-  ResultGraph->markIncompleteNodes(DSGraph::IgnoreFormalArgs |
-                                   DSGraph::IgnoreGlobals);
+  ResultGraph->markIncompleteNodes(DSGraph::IgnoreGlobals |
+                                   DSGraph::MarkFormalArgs);
 
   // Remove any nodes that are dead after all of the merging we have done...
   // FIXME: We should be able to disable the globals graph for steens!