--- a/nashorn/src/jdk.scripting.nashorn/share/classes/jdk/nashorn/api/tree/CompilationUnitTreeImpl.java	Wed Jul 05 21:57:11 2017 +0200
+++ b/nashorn/src/jdk.scripting.nashorn/share/classes/jdk/nashorn/api/tree/CompilationUnitTreeImpl.java	Tue Jul 12 21:18:13 2016 +0530
@@ -32,13 +32,18 @@
     implements CompilationUnitTree {
     private final FunctionNode funcNode;
     private final List<? extends Tree> elements;
+    private final ModuleTree module;
 
     CompilationUnitTreeImpl(final FunctionNode node,
-            final List<? extends Tree> elements) {
+            final List<? extends Tree> elements,
+            final ModuleTree module) {
         super(node);
         this.funcNode = node;
-        assert funcNode.getKind() == FunctionNode.Kind.SCRIPT : "script function expected";
+        assert funcNode.getKind() == FunctionNode.Kind.SCRIPT ||
+                funcNode.getKind() == FunctionNode.Kind.MODULE :
+                "script or module function expected";
         this.elements = elements;
+        this.module = module;
     }
 
     @Override
@@ -67,6 +72,11 @@
     }
 
     @Override
+    public ModuleTree getModule() {
+        return module;
+    }
+
+    @Override
     public <R,D> R accept(final TreeVisitor<R,D> visitor, final D data) {
         return visitor.visitCompilationUnit(this, data);
     }