Add ability to stop scripts

recaf-core/src/main/java/software/coley/recaf/services/plugin/CdiClassAllocator.java

@@ -6,8 +6,10 @@
 import jakarta.enterprise.inject.spi.*;
 import jakarta.inject.Inject;
 
-import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Allocator instance that ties into the CDI container.
@@ -16,7 +18,8 @@
  */
 @ApplicationScoped
 public class CdiClassAllocator implements ClassAllocator {
-	private final Map<Class<?>, Bean<?>> classBeanMap = new IdentityHashMap<>();
+	private final Map<Class<?>, Bean<?>> classBeanMap = new WeakHashMap<>();
+	private final Lock lock = new ReentrantLock();
 	private final BeanManager beanManager;
 
 	@Inject
@@ -28,20 +31,27 @@ public CdiClassAllocator(@Nonnull BeanManager beanManager) {
 	@Override
 	@SuppressWarnings("unchecked")
 	public <T> T instance(@Nonnull Class<T> cls) throws AllocationException {
+		lock.lock();
 		try {
 			// Create bean
 			Bean<T> bean = (Bean<T>) classBeanMap.computeIfAbsent(cls, c -> {
-				AnnotatedType<T> annotatedClass = beanManager.createAnnotatedType(cls);
+				// TODO something here is bugged
+				//  bean.create(creationalContext).getClass().getClassLoader() == cls.getClassLoader()
+				//   - Evaluates to false
+				//   - In AnnotatedTypeIdentifier#forBackedAnnotatedType the ClassLoader is not considered, just the class's name
+				AnnotatedType<T> annotatedClass = beanManager.createAnnotatedType((Class<T>) c);
 				BeanAttributes<T> attributes = beanManager.createBeanAttributes(annotatedClass);
 				InjectionTargetFactory<T> factory = beanManager.getInjectionTargetFactory(annotatedClass);
-				return beanManager.createBean(attributes, cls, factory);
+				return beanManager.createBean(attributes, (Class<T>) c, factory);
 			});
 			CreationalContext<T> creationalContext = beanManager.createCreationalContext(bean);
 
 			// Allocate instance of bean
 			return bean.create(creationalContext);
 		} catch (Throwable t) {
 			throw new AllocationException(cls, t);
+		} finally {
+			lock.unlock();
 		}
 	}
 }

recaf-core/src/main/java/software/coley/recaf/services/script/CancellationSingleton.java

@@ -0,0 +1,44 @@
+package software.coley.recaf.services.script;
+
+import software.coley.recaf.util.CancelSignal;
+
+/**
+ * Cancellation singleton.
+ * <p>
+ * Injected into generated scripts.
+ * <b>Do not use in normal code</b>.
+ *
+ * @author xDark
+ * @see InsertCancelSignalVisitor
+ * @see GenerateResult#requestStop()
+ */
+public final class CancellationSingleton {
+	private static volatile boolean shouldStop;
+
+	private CancellationSingleton() {}
+
+	/**
+	 * Schedules cancellation of the script.
+	 */
+	public static void stop() {
+		shouldStop = true;
+	}
+
+	/**
+	 * Clears any prior cancellation request.
+	 */
+	public static void reset() {
+		shouldStop = false;
+	}
+
+	/**
+	 * Polls for cancellation.
+	 *
+	 * @throws CancelSignal
+	 * 		If cancellation has been requested.
+	 */
+	public static void poll() {
+		if (shouldStop)
+			throw CancelSignal.get();
+	}
+}

recaf-core/src/main/java/software/coley/recaf/services/script/GenerateResult.java

@@ -4,6 +4,7 @@
 import jakarta.annotation.Nullable;
 import software.coley.recaf.services.compile.CompilerDiagnostic;
 
+import java.lang.reflect.InvocationTargetException;
 import java.util.List;
 
 /**
@@ -23,4 +24,42 @@ public record GenerateResult(@Nullable Class<?> cls, @Nonnull List<CompilerDiagn
 	public boolean wasSuccess() {
 		return cls != null;
 	}
-}
+
+	/**
+	 * Attempts to stop the script. If the generation failed, this method will do nothing.
+	 * <p>
+	 * Cancellation is scoped to the generated class loader that owns {@link #cls()}.
+	 * Callers that wish to stop running scripts must use {@link ScriptEngine#run(GenerateResult)}
+	 * and track the returned instance for stopping.
+	 *
+	 * @throws IllegalStateException
+	 * 		If something went wrong.
+	 */
+	public void requestStop() {
+		invokeCancellationSingleton("stop");
+	}
+
+	/**
+	 * Clears any prior request to stop the script. If the generation failed, this method will do nothing.
+	 *
+	 * @throws IllegalStateException
+	 * 		If something went wrong.
+	 */
+	public void resetStop() {
+		invokeCancellationSingleton("reset");
+	}
+
+	private void invokeCancellationSingleton(@Nonnull String methodName) {
+		Class<?> cls = this.cls;
+		if (cls == null)
+			return;
+		try {
+			Class<?> cancellationSingleton = cls.getClassLoader().loadClass(CancellationSingleton.class.getName());
+			cancellationSingleton.getDeclaredMethod(methodName).invoke(null);
+		} catch (InvocationTargetException ex) {
+			throw new IllegalStateException(ex.getTargetException());
+		} catch (Exception ex) {
+			throw new IllegalStateException(ex);
+		}
+	}
+}