Issue 154: Add null-check to @Cleanup

6 files changed, 67 insertions, 13 deletions

diff --git a/src/core/lombok/eclipse/handlers/HandleCleanup.java b/src/core/lombok/eclipse/handlers/HandleCleanup.java
index d296e96b..33ab7fb9 100644
--- a/src/core/lombok/eclipse/handlers/HandleCleanup.java
+++ b/src/core/lombok/eclipse/handlers/HandleCleanup.java
@@ -37,9 +37,13 @@ import org.eclipse.jdt.internal.compiler.ast.Assignment;
 import org.eclipse.jdt.internal.compiler.ast.Block;
 import org.eclipse.jdt.internal.compiler.ast.CaseStatement;
 import org.eclipse.jdt.internal.compiler.ast.CastExpression;
+import org.eclipse.jdt.internal.compiler.ast.EqualExpression;
+import org.eclipse.jdt.internal.compiler.ast.IfStatement;
 import org.eclipse.jdt.internal.compiler.ast.LocalDeclaration;
 import org.eclipse.jdt.internal.compiler.ast.MemberValuePair;
 import org.eclipse.jdt.internal.compiler.ast.MessageSend;
+import org.eclipse.jdt.internal.compiler.ast.NullLiteral;
+import org.eclipse.jdt.internal.compiler.ast.OperatorIds;
 import org.eclipse.jdt.internal.compiler.ast.SingleNameReference;
 import org.eclipse.jdt.internal.compiler.ast.Statement;
 import org.eclipse.jdt.internal.compiler.ast.SwitchStatement;
@@ -157,7 +161,29 @@ public class HandleCleanup implements EclipseAnnotationHandler<Cleanup> {
 		}
 		unsafeClose.nameSourcePosition = nameSourcePosition;
 		unsafeClose.selector = cleanupName.toCharArray();
-		finallyBlock[0] = unsafeClose;
+		
+		
+		int pS = ast.sourceStart, pE = ast.sourceEnd;
+		long p = (long)pS << 32 | pE;
+
+		SingleNameReference varName = new SingleNameReference(decl.name, p);
+		Eclipse.setGeneratedBy(varName, ast);
+		NullLiteral nullLiteral = new NullLiteral(pS, pE);
+		Eclipse.setGeneratedBy(nullLiteral, ast);
+		EqualExpression equalExpression = new EqualExpression(varName, nullLiteral, OperatorIds.NOT_EQUAL);
+		equalExpression.sourceStart = pS; equalExpression.sourceEnd = pE;
+		Eclipse.setGeneratedBy(equalExpression, ast);
+		
+		Block closeBlock = new Block(0);
+		closeBlock.statements = new Statement[1];
+		closeBlock.statements[0] = unsafeClose;
+		Eclipse.setGeneratedBy(closeBlock, ast);
+		IfStatement ifStatement = new IfStatement(equalExpression, closeBlock, 0, 0);
+		Eclipse.setGeneratedBy(ifStatement, ast);
+
+		
+		
+		finallyBlock[0] = ifStatement;
 		tryStatement.finallyBlock = new Block(0);
 		Eclipse.setGeneratedBy(tryStatement.finallyBlock, ast);
 		tryStatement.finallyBlock.statements = finallyBlock;
diff --git a/src/core/lombok/javac/handlers/HandleCleanup.java b/src/core/lombok/javac/handlers/HandleCleanup.java
index 2c89d9ad..779dd3ea 100644
--- a/src/core/lombok/javac/handlers/HandleCleanup.java
+++ b/src/core/lombok/javac/handlers/HandleCleanup.java
@@ -30,10 +30,12 @@ import lombok.javac.JavacNode;
 
 import org.mangosdk.spi.ProviderFor;
 
+import com.sun.tools.javac.code.TypeTags;
 import com.sun.tools.javac.tree.JCTree;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.tree.JCTree.JCAnnotation;
 import com.sun.tools.javac.tree.JCTree.JCAssign;
+import com.sun.tools.javac.tree.JCTree.JCBinary;
 import com.sun.tools.javac.tree.JCTree.JCBlock;
 import com.sun.tools.javac.tree.JCTree.JCCase;
 import com.sun.tools.javac.tree.JCTree.JCCatch;
@@ -41,6 +43,7 @@ import com.sun.tools.javac.tree.JCTree.JCExpression;
 import com.sun.tools.javac.tree.JCTree.JCExpressionStatement;
 import com.sun.tools.javac.tree.JCTree.JCFieldAccess;
 import com.sun.tools.javac.tree.JCTree.JCIdent;
+import com.sun.tools.javac.tree.JCTree.JCIf;
 import com.sun.tools.javac.tree.JCTree.JCMethodDecl;
 import com.sun.tools.javac.tree.JCTree.JCStatement;
 import com.sun.tools.javac.tree.JCTree.JCTypeCast;
@@ -108,11 +111,16 @@ public class HandleCleanup implements JavacAnnotationHandler<Cleanup> {
 		doAssignmentCheck(annotationNode, tryBlock, decl.name);
 		
 		TreeMaker maker = annotationNode.getTreeMaker();
-		JCFieldAccess cleanupCall = maker.Select(maker.Ident(decl.name), annotationNode.toName(cleanupName));
-		List<JCStatement> finalizerBlock = List.<JCStatement>of(maker.Exec(
-				maker.Apply(List.<JCExpression>nil(), cleanupCall, List.<JCExpression>nil())));
+		JCFieldAccess cleanupMethod = maker.Select(maker.Ident(decl.name), annotationNode.toName(cleanupName));
+		List<JCStatement> cleanupCall = List.<JCStatement>of(maker.Exec(
+				maker.Apply(List.<JCExpression>nil(), cleanupMethod, List.<JCExpression>nil())));
+		
+		JCBinary isNull = maker.Binary(JCTree.NE, maker.Ident(decl.name), maker.Literal(TypeTags.BOT, null));
+		
+		JCIf ifNotNullCleanup = maker.If(isNull, maker.Block(0, cleanupCall), null);
+		
+		JCBlock finalizer = maker.Block(0, List.<JCStatement>of(ifNotNullCleanup));
 		
-		JCBlock finalizer = maker.Block(0, finalizerBlock);
 		newStatements = newStatements.append(maker.Try(maker.Block(0, tryBlock), List.<JCCatch>nil(), finalizer));
 		
 		if (blockNode instanceof JCBlock) {
diff --git a/test/transform/resource/after-delombok/CleanupName.java b/test/transform/resource/after-delombok/CleanupName.java
index cd29eb68..0201008e 100644
--- a/test/transform/resource/after-delombok/CleanupName.java
+++ b/test/transform/resource/after-delombok/CleanupName.java
@@ -4,7 +4,9 @@ class CleanupName {
 		try {
 			System.out.println(o);
 		} finally {
-			o.toString();
+			if (o != null) {
+				o.toString();
+			}
 		}
 	}
 	void test2() {
@@ -12,7 +14,9 @@ class CleanupName {
 		try {
 			System.out.println(o);
 		} finally {
-			o.toString();
+			if (o != null) {
+				o.toString();
+			}
 		}
 	}
 }
diff --git a/test/transform/resource/after-delombok/CleanupPlain.java b/test/transform/resource/after-delombok/CleanupPlain.java
index 35d51543..1a19442f 100644
--- a/test/transform/resource/after-delombok/CleanupPlain.java
+++ b/test/transform/resource/after-delombok/CleanupPlain.java
@@ -9,10 +9,14 @@ class CleanupPlain {
 					out.flush();
 				}
 			} finally {
-				out.close();
+				if (out != null) {
+					out.close();
+				}
 			}
 		} finally {
-			in.close();
+			if (in != null) {
+				in.close();
+			}
 		}
 	}
 }
diff --git a/test/transform/resource/after-ecj/CleanupName.java b/test/transform/resource/after-ecj/CleanupName.java
index 944a81e1..4011721b 100644
--- a/test/transform/resource/after-ecj/CleanupName.java
+++ b/test/transform/resource/after-ecj/CleanupName.java
@@ -10,7 +10,10 @@ class CleanupName {
       }
     finally
       {
-        o.toString();
+        if ((o != null))
+            {
+              o.toString();
+            }
       }
   }
   void test2() {
@@ -21,7 +24,10 @@ class CleanupName {
       }
     finally
       {
-        o.toString();
+        if ((o != null))
+            {
+              o.toString();
+            }
       }
   }
 }
\ No newline at end of file
diff --git a/test/transform/resource/after-ecj/CleanupPlain.java b/test/transform/resource/after-ecj/CleanupPlain.java
index 6754b4f6..fe8d7571 100644
--- a/test/transform/resource/after-ecj/CleanupPlain.java
+++ b/test/transform/resource/after-ecj/CleanupPlain.java
@@ -18,12 +18,18 @@ class CleanupPlain {
           }
         finally
           {
-            out.close();
+            if ((out != null))
+                {
+                  out.close();
+                }
           }
       }
     finally
       {
-        in.close();
+        if ((in != null))
+            {
+              in.close();
+            }
       }
   }
 }
\ No newline at end of file