package com.redhat.ceylon.eclipse.core.model;

import java.lang.ref.WeakReference;

import com.redhat.ceylon.compiler.typechecker.context.PhasedUnit;
import com.redhat.ceylon.eclipse.core.typechecker.IdePhasedUnit;

public abstract class CeylonUnit extends IdeUnit {
    
    public CeylonUnit() {
        phasedUnitRef = null;
    }
    
    protected WeakReference<IdePhasedUnit> phasedUnitRef;
    
    final protected void createPhasedUnitRef(IdePhasedUnit phasedUnit) {
        phasedUnitRef = new WeakReference<IdePhasedUnit>(phasedUnit);
    }
    
    protected abstract void setPhasedUnitIfNecessary();
    
    public IdePhasedUnit getPhasedUnit() {
        setPhasedUnitIfNecessary();
        return phasedUnitRef.get();
    }
}

@@ -143,16 +143,73 @@
 edUnit;%0A
+import com.redhat.ceylon.compiler.typechecker.tree.Tree;%0A
 import c
@@ -788,10 +788,226 @@
 ;%0A    %7D%0A
+    %0A    public Tree.CompilationUnit getCompilationUnit() %7B%0A        IdePhasedUnit pu = getPhasedUnit();%0A        if (pu == null) %7B%0A            return null;%0A        %7D%0A        return pu.getCompilationUnit();%0A    %7D%0A    %0A
 %7D%0A

Add a CompilationUnit getter on the model Unit (to avoid testing for null each time)