From 5ed634a60e09044e03928fef7f9322e647273ab9 Mon Sep 17 00:00:00 2001
From: Knut Wannheden <knut@moderne.io>
Date: Mon, 7 Oct 2024 15:03:29 +0200
Subject: [PATCH] Implement `JS.CompilationUnit#getWeight()`

The standard `SourceFile` implementation is incompatible with GraalVM, as it relies on dynamic class generation via `TreeVisitorAdapter` (and Gizmo).
---
 .../org/openrewrite/javascript/tree/JS.java   | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/src/main/java/org/openrewrite/javascript/tree/JS.java b/src/main/java/org/openrewrite/javascript/tree/JS.java
index f12cacc0..4c17b4e3 100644
--- 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;