more checking

git-svn-id: https://wala.svn.sourceforge.net/svnroot/wala/trunk@3937 f5eafffb-2e1d-0410-98e4-8ec43c5233c4

com.ibm.wala.cast/source/java/com/ibm/wala/cast/tree/impl/CAstControlFlowRecorder.java

@@ -72,6 +72,10 @@ public class CAstControlFlowRecorder implements CAstControlFlowMap {
       return (o instanceof Key) && from == ((Key) o).from
           && ((label == null) ? ((Key) o).label == null : label.equals(((Key) o).label));
     }
+    
+    public String toString() {
+      return "<key " + label + " : " + from + ">";
+    }
   }
 
   public CAstControlFlowRecorder(CAstSourcePositionMap src) {
@@ -80,10 +84,13 @@ public class CAstControlFlowRecorder implements CAstControlFlowMap {
   }
 
   public CAstNode getTarget(CAstNode from, Object label) {
+    assert CAstToNode.get(from) != null;
     Key key = new Key(label, CAstToNode.get(from));
-    if (table.containsKey(key))
-      return (CAstNode) nodeToCAst.get(table.get(key));
-    else
+    if (table.containsKey(key)) {
+      Object target = table.get(key);
+      assert nodeToCAst.containsKey(target);
+      return (CAstNode) nodeToCAst.get(target);
+    } else
       return null;
   }
 
@@ -105,8 +112,10 @@ public class CAstControlFlowRecorder implements CAstControlFlowMap {
 
   public Collection<CAstNode> getMappedNodes() {
     Set<CAstNode> nodes = new LinkedHashSet<CAstNode>();
-    for (Iterator keys = table.keySet().iterator(); keys.hasNext();) {
-      nodes.add((CAstNode) nodeToCAst.get(((Key) keys.next()).from));
+    for (Iterator<Key> keys = table.keySet().iterator(); keys.hasNext();) {
+      Key key = keys.next();
+      nodes.add((CAstNode) nodeToCAst.get(key.from));
+      nodes.add((CAstNode) nodeToCAst.get(table.get(key)));
     }
 
     return nodes;
@@ -119,6 +128,9 @@ public class CAstControlFlowRecorder implements CAstControlFlowMap {
    * this add call.
    */
   public void add(Object from, Object to, Object label) {
+    assert from != null;
+    assert to != null;
+    
     table.put(new Key(label, from), to);
 
     Set<Object> ls = labelMap.get(from);
@@ -134,7 +146,7 @@ public class CAstControlFlowRecorder implements CAstControlFlowMap {
 
   /**
    * Establish a mapping between some object `node' and the ast node `ast'.
-   * Objects used a endpoints in a control flow edge must be mapped to ast nodes
+   * Objects used as endpoints in a control flow edge must be mapped to ast nodes
    * using this call.
    */
   public void map(Object node, CAstNode ast) {

com.ibm.wala.cast/source/java/com/ibm/wala/cast/tree/impl/CAstImpl.java

@@ -66,7 +66,7 @@ public class CAstImpl implements CAst {
     }
 
     public String toString() {
-      return super.toString() + ":" + CAstPrinter.print(this);
+      return System.identityHashCode(this) + ":" + CAstPrinter.print(this);
     }
 
     public int hashCode() {

com.ibm.wala.cast/source/java/com/ibm/wala/cast/tree/impl/CAstRewriter.java

@@ -85,6 +85,7 @@ public abstract class CAstRewriter<C extends CAstRewriter.RewriteContext<K>, K e
 
   private CAstControlFlowMap copyFlow(Map<Pair, CAstNode> nodeMap, CAstControlFlowMap orig, CAstSourcePositionMap newSrc) {
     Set<CAstNode> mappedOutsideNodes = HashSetFactory.make(1);
+    Set<CAstNode> allNewTargetNodes = HashSetFactory.make(1);
     CAstControlFlowRecorder newMap = new CAstControlFlowRecorder(newSrc);
     Collection oldSources = orig.getMappedNodes();
 
@@ -112,6 +113,7 @@ public abstract class CAstRewriter<C extends CAstRewriter.RewriteContext<K>, K e
         while (LS.hasNext()) {
           Object label = LS.next();
           CAstNode oldTarget = orig.getTarget(oldSource, label);
+          assert oldTarget != null;
 
           if (DEBUG) {
             System.err.println(("old: " + label + " --> " + CAstPrinter.print(oldTarget)));
@@ -132,9 +134,11 @@ public abstract class CAstRewriter<C extends CAstRewriter.RewriteContext<K>, K e
           if (nodeMap.containsKey(targetKey)) {
             newTarget = (CAstNode) nodeMap.get(targetKey);
             newMap.add(newSource, newTarget, label);
+            allNewTargetNodes.add(newTarget);
 
           } else {
             newTarget = flowOutTo(nodeMap, oldSource, label, oldTarget, orig, newSrc);
+            allNewTargetNodes.add(newTarget);
             newMap.add(newSource, newTarget, label);
             if (newTarget != CAstControlFlowMap.EXCEPTION_TO_EXIT && !mappedOutsideNodes.contains(newTarget)) {
               mappedOutsideNodes.add(newTarget);
@@ -150,6 +154,11 @@ public abstract class CAstRewriter<C extends CAstRewriter.RewriteContext<K>, K e
       }
     }
 
+    allNewTargetNodes.removeAll(newMap.getMappedNodes());
+    for(CAstNode newTarget : allNewTargetNodes) {
+      newMap.map(newTarget, newTarget);
+    }
+    
     return newMap;
   }
 
@@ -269,9 +278,7 @@ public abstract class CAstRewriter<C extends CAstRewriter.RewriteContext<K>, K e
   }
 
   public CAstEntity rewrite(final CAstEntity root) {
-
-    System.err.println(("Rewriting " + root.getName()));
-
+    
     if (root.getAST() != null) {
       final Rewrite rewrite = rewrite(root.getAST(), root.getControlFlow(), root.getSourceMap(), root.getNodeTypeMap(), root
           .getAllScopedEntities());