Issue 154: Add null-check to @Cleanup

2 files changed, 39 insertions, 5 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) {