Implement JS.CompilationUnit#getWeight()

The standard `SourceFile` implementation is incompatible with GraalVM, as it relies on dynamic class generation via `TreeVisitorAdapter` (and Gizmo).
openrewrite · Oct 7, 2024 · 5ed634a · 5ed634a
1 parent aadfb9b
commit 5ed634a
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/src/main/java/org/openrewrite/javascript/tree/JS.java b/src/main/java/org/openrewrite/javascript/tree/JS.java
@@ -21,6 +21,7 @@
 import org.jspecify.annotations.Nullable;
 import org.openrewrite.*;
 import org.openrewrite.java.JavaPrinter;
+import org.openrewrite.java.JavaTypeVisitor;
 import org.openrewrite.java.internal.TypesInUse;
 import org.openrewrite.java.service.ImportService;
 import org.openrewrite.java.tree.*;
@@ -38,6 +39,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 import static java.util.Collections.singletonList;
@@ -228,6 +231,37 @@ public <S, T extends S> T service(Class<S> service) {
             }
         }
 
+        @Transient
+        @Override
+        public long getWeight(Predicate<Object> uniqueIdentity) {
+            AtomicInteger n = new AtomicInteger();
+            new JavaScriptVisitor<AtomicInteger>() {
+                final JavaTypeVisitor<AtomicInteger> typeVisitor = new JavaTypeVisitor<AtomicInteger>() {
+                    @Override
+                    public JavaType visit(@Nullable JavaType javaType, AtomicInteger n) {
+                        if (javaType != null && uniqueIdentity.test(javaType)) {
+                            n.incrementAndGet();
+                            return super.visit(javaType, n);
+                        }
+                        //noinspection ConstantConditions
+                        return javaType;
+                    }
+                };
+
+                @Override
+                public J preVisit(J tree, AtomicInteger n) {
+                    n.incrementAndGet();
+                    return tree;
+                }
+
+                @Override
+                public JavaType visitType(@Nullable JavaType javaType, AtomicInteger n) {
+                    return typeVisitor.visit(javaType, n);
+                }
+            }.visit(this, n);
+            return n.get();
+        }
+
         @Override
         public Padding getPadding() {
             Padding p;