Merge pull request #1314 from NASA-AMMOS/feature--streamline-resource…

…-profiling

Add more utilities for profiling resources

contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/MutableResource.java

@@ -12,6 +12,8 @@
 import static gov.nasa.jpl.aerie.contrib.streamline.core.CellRefV2.autoEffects;
 import static gov.nasa.jpl.aerie.contrib.streamline.core.monads.DynamicsMonad.pure;
 import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Naming.*;
+import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling.profile;
+import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling.profileEffects;
 import static java.util.stream.Collectors.joining;
 
 /**
@@ -62,7 +64,10 @@ private void augmentEffectName(DynamicsEffect<D> effect) {
       }
     };
     if (MutableResourceFlags.DETECT_BUSY_CELLS) {
-      result = Profiling.profileEffects(result);
+      result = profileEffects(result);
+    }
+    if (MutableResourceFlags.PROFILE_GET_DYNAMICS) {
+      result = profile(result);
     }
     return result;
   }
@@ -79,12 +84,12 @@ static <D extends Dynamics<?, D>> void set(MutableResource<D> resource, Expiring
    * Turn on busy cell detection.
    *
    * <p>
-   *     Calling this method once before constructing your model will profile effects on every cell.
+   *     Calling this method once before constructing your model will profile effects on every resource.
    *     Profiling effects may be compute and/or memory intensive, and should not be used in production.
    * </p>
    * <p>
-   *     If only a few cells are suspect, you can also call {@link Profiling#profileEffects}
-   *     directly on just those cells, rather than profiling every cell.
+   *     If only a few resources are suspect, you can also call {@link Profiling#profileEffects}
+   *     directly on just those resource, rather than profiling every resource.
    * </p>
    * <p>
    *     Call {@link Profiling#dump()} to see results.
@@ -93,6 +98,27 @@ static <D extends Dynamics<?, D>> void set(MutableResource<D> resource, Expiring
   static void detectBusyCells() {
     MutableResourceFlags.DETECT_BUSY_CELLS = true;
   }
+
+  /**
+   * Turn on profiling for all {@link MutableResource}s created by {@link MutableResource#resource}.
+   * Also implies {@link MutableResource#detectBusyCells()}.
+   *
+   * <p>
+   *     Calling this method once before constructing your model will profile virtually every {@link MutableResource}.
+   *     Profiling may be compute and/or memory intensive, and should not be used in production.
+   * </p>
+   * <p>
+   *     If only a few resources are suspect, you can also call {@link Profiling#profile}
+   *     directly on just those resource, rather than profiling every resource.
+   * </p>
+   * <p>
+   *     Call {@link Profiling#dump()} to see results.
+   * </p>
+   */
+  static void profileAllResources() {
+    MutableResourceFlags.PROFILE_GET_DYNAMICS = true;
+    detectBusyCells();
+  }
 }
 
 /**
@@ -102,4 +128,5 @@ static void detectBusyCells() {
  */
 final class MutableResourceFlags {
   public static boolean DETECT_BUSY_CELLS = false;
+  public static boolean PROFILE_GET_DYNAMICS = false;
 }

contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/Resource.java

@@ -1,8 +1,31 @@
 package gov.nasa.jpl.aerie.contrib.streamline.core;
 
+import gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad;
+import gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling;
+
 /**
  * A function returning a fully-wrapped dynamics,
  * and the primary way models track state and report results.
  */
 public interface Resource<D> extends ThinResource<ErrorCatching<Expiring<D>>> {
+    /**
+     * Turn on profiling for all resources derived through {@link ResourceMonad}
+     * or created by {@link MutableResource#resource}.
+     *
+     * <p>
+     *     Calling this method once before constructing your model will profile virtually every resource.
+     *     Profiling may be compute and/or memory intensive, and should not be used in production.
+     * </p>
+     * <p>
+     *     If only a few resources are suspect, you can also call {@link Profiling#profile}
+     *     directly on just those resource, rather than profiling every resource.
+     * </p>
+     * <p>
+     *     Call {@link Profiling#dump()} to see results.
+     * </p>
+     */
+    static void profileAllResources() {
+        ResourceMonad.profileAllResources();
+        MutableResource.profileAllResources();
+    }
 }

contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/monads/ResourceMonad.java

@@ -4,13 +4,15 @@
 import gov.nasa.jpl.aerie.contrib.streamline.core.Expiring;
 import gov.nasa.jpl.aerie.contrib.streamline.core.Resource;
 import gov.nasa.jpl.aerie.contrib.streamline.core.ThinResource;
+import gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling;
 import gov.nasa.jpl.aerie.contrib.streamline.utils.*;
 import org.apache.commons.lang3.function.TriFunction;
 
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Dependencies.addDependency;
+import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling.profile;
 import static gov.nasa.jpl.aerie.contrib.streamline.utils.FunctionalUtils.curry;
 
 /**
@@ -21,14 +23,37 @@
 public final class ResourceMonad {
   private ResourceMonad() {}
 
+  private static boolean profileAllResources = false;
+  /**
+   * Turn on profiling for all getDynamics calls on {@link Resource}s derived through {@link ResourceMonad}.
+   *
+   * <p>
+   *     Calling this method once before constructing your model will profile getDynamics on every derived resource.
+   *     Profiling may be compute and/or memory intensive, and should not be used in production.
+   * </p>
+   * <p>
+   *     If only a few cells are suspect, you can also call {@link Profiling#profile}
+   *     directly on just those resource, rather than profiling every resource.
+   * </p>
+   * <p>
+   *     Call {@link Profiling#dump()} to see results.
+   * </p>
+   */
+  public static void profileAllResources() {
+    profileAllResources = true;
+  }
+
   public static <A> Resource<A> pure(A a) {
-    return ThinResourceMonad.pure(DynamicsMonad.pure(a))::getDynamics;
+    Resource<A> result = ThinResourceMonad.pure(DynamicsMonad.pure(a))::getDynamics;
+    if (profileAllResources) result = profile(result);
+    return result;
   }
 
   public static <A, B> Resource<B> apply(Resource<A> a, Resource<Function<A, B>> f) {
     Resource<B> result = ThinResourceMonad.apply(a, ThinResourceMonad.map(f, DynamicsMonad::apply))::getDynamics;
     addDependency(result, a);
     addDependency(result, f);
+    if (profileAllResources) result = profile(result);
     return result;
   }
 
@@ -43,6 +68,7 @@ public static <A> Resource<A> join(Resource<Resource<A>> a) {
     // The ::getDynamics at the end up-converts back to Resource, from ThinResource
     Resource<A> result = ThinResourceMonad.map(ThinResourceMonad.join(ThinResourceMonad.map(a$, ResourceMonad::distribute)), DynamicsMonad::join)::getDynamics;
     addDependency(result, a);
+    if (profileAllResources) result = profile(result);
     return result;
   }
 

contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/debugging/Naming.java

@@ -4,6 +4,7 @@
 
 import java.lang.ref.WeakReference;
 import java.util.*;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -25,9 +26,19 @@ public final class Naming {
   private Naming() {}
 
   // Use a WeakHashMap so that naming a thing doesn't prevent it from being garbage-collected.
-  private static final WeakHashMap<Object, Supplier<Optional<String>>> NAMES = new WeakHashMap<>();
-  // Way to inject a temporary "anonymous" name, so derived names still work even when not all args are named.
-  private static final MutableObject<Optional<String>> anonymousName = new MutableObject<>(Optional.empty());
+  private static final WeakHashMap<Object, Function<NamingContext, Optional<String>>> NAMES = new WeakHashMap<>();
+
+  private record NamingContext(Set<Object> visited, Optional<String> anonymousName) {
+    NamingContext visit(Object thing) {
+      var newVisited = new HashSet<>(visited);
+      newVisited.add(thing);
+      return new NamingContext(newVisited, anonymousName);
+    }
+
+    public NamingContext(String anonymousName) {
+      this(Set.of(), Optional.ofNullable(anonymousName));
+    }
+  }
 
   /**
    * Register a name for thing, as a function of args' names.
@@ -36,15 +47,12 @@ private Naming() {}
   public static <T> T name(T thing, String nameFormat, Object... args) {
     // Only capture weak references to arguments, so we don't leak memory
     var args$ = Arrays.stream(args).map(WeakReference::new).toArray(WeakReference[]::new);
-    NAMES.put(thing, () -> {
+    NAMES.put(thing, context -> {
       Object[] argNames = new Object[args$.length];
       for (int i = 0; i < args$.length; ++i) {
-        // Try to resolve the argument name by first looking up and using its registered name,
-        // or by falling back to the anonymous name.
         var argName$ = Optional.ofNullable(args$[i].get())
-                .flatMap(Naming::getName)
-                .or(anonymousName::getValue);
-        if (argName$.isEmpty()) return Optional.empty();
+                .flatMap(argRef -> getName(argRef, context));
+        if (argName$.isEmpty()) return context.anonymousName();
         argNames[i] = argName$.get();
       }
       return Optional.of(nameFormat.formatted(argNames));
@@ -58,19 +66,22 @@ public static <T> T name(T thing, String nameFormat, Object... args) {
    * returns empty.
    */
   public static Optional<String> getName(Object thing) {
-    return Optional.ofNullable(NAMES.get(thing)).flatMap(Supplier::get).or(anonymousName::getValue);
+    return getName(thing, new NamingContext(null));
   }
 
   /**
    * Get the name for thing.
    * Use anonymousName for anything without a name instead of returning empty.
    */
   public static String getName(Object thing, String anonymousName) {
-    Naming.anonymousName.setValue(Optional.of(anonymousName));
-    var result = getName(thing);
-    Naming.anonymousName.setValue(Optional.empty());
-    // This will never throw, because anonymous name will guarantee that some name is found.
-    return result.orElseThrow();
+    // This expression never throws, because context always has a name available.
+    return getName(thing, new NamingContext(anonymousName)).orElseThrow();
+  }
+
+  private static Optional<String> getName(Object thing, NamingContext context) {
+    return context.visited.contains(thing)
+            ? context.anonymousName
+            : NAMES.getOrDefault(thing, NamingContext::anonymousName).apply(context.visit(thing));
   }
 
   public static String argsFormat(Collection<?> collection) {