Skip to content

Memoize safety evaluator #2812

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xanderbailey wants to merge 5 commits into
base: develop
Choose a base branch
from
Open

Memoize safety evaluator #2812

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a6cbba2
Memoize safety evaluator
xanderbailey Feb 23, 2026
7976a63
docs
xanderbailey Feb 23, 2026
d8ee4e0
fixup docs
xanderbailey Feb 23, 2026
d55ef72
Add tests
xanderbailey Feb 23, 2026
b9870b6
format
xanderbailey Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 44 additions & 3 deletions conjure-java-core/src/main/java/com/palantir/conjure/java/types/SafetyEvaluator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import com.palantir.conjure.spec.TypeName;
import com.palantir.conjure.spec.UnionDefinition;
import com.palantir.logsafe.Preconditions;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -65,6 +66,10 @@ public final class SafetyEvaluator {

private final Map<TypeName, TypeDefinition> definitionMap;

// Memoization cache shared across all evaluate() calls on this instance. Avoids redundant recursive
// traversals of the type graph, which otherwise dominate generation time for large definitions.
private final Map<TypeName, Optional<LogSafety>> cache = new HashMap<>();

public SafetyEvaluator(ConjureDefinition definition) {
this(TypeFunctions.toTypesMap(definition));
}
Expand All @@ -75,12 +80,12 @@ public SafetyEvaluator(Map<TypeName, TypeDefinition> definitionMap) {

public Optional<LogSafety> evaluate(TypeDefinition def) {
return Preconditions.checkNotNull(def, "TypeDefinition is required")
.accept(new TypeDefinitionSafetyVisitor(definitionMap, new HashSet<>()));
.accept(new TypeDefinitionSafetyVisitor(definitionMap, cache, new HashSet<>()));
}

public Optional<LogSafety> evaluate(Type type) {
return Preconditions.checkNotNull(type, "TypeDefinition is required")
.accept(new TypeDefinitionSafetyVisitor(definitionMap, new HashSet<>()).fieldVisitor);
.accept(new TypeDefinitionSafetyVisitor(definitionMap, cache, new HashSet<>()).fieldVisitor);
}

public Optional<LogSafety> evaluate(Type type, Optional<LogSafety> declaredSafety) {
Expand Down Expand Up @@ -124,10 +129,19 @@ public Optional<LogSafety> getUsageTimeSafety(FieldDefinition field) {
}

private static final class TypeDefinitionSafetyVisitor implements TypeDefinition.Visitor<Optional<LogSafety>> {
private final Map<TypeName, Optional<LogSafety>> cache;
private final Set<TypeName> inProgress;
private final Type.Visitor<Optional<LogSafety>> fieldVisitor;

private TypeDefinitionSafetyVisitor(Map<TypeName, TypeDefinition> definitionMap, Set<TypeName> inProgress) {
// Tracks whether cycle-breaking (the SAFE fallback for back-edges) was used anywhere
// in the current evaluation subtree. Used to decide whether a result is safe to cache.
private boolean encounteredCycle;

private TypeDefinitionSafetyVisitor(
Map<TypeName, TypeDefinition> definitionMap,
Map<TypeName, Optional<LogSafety>> cache,
Set<TypeName> inProgress) {
this.cache = cache;
this.inProgress = inProgress;
this.fieldVisitor = new FieldSafetyVisitor(definitionMap, this);
}
Expand Down Expand Up @@ -170,15 +184,42 @@ public Optional<LogSafety> visitUnknown(String unknownType) {
}

private Optional<LogSafety> with(TypeName typeName, Supplier<Optional<LogSafety>> task) {
// Return memoized result if this type has already been fully evaluated.
// Note: cache values are Optional<LogSafety> which may be Optional.empty(),
// so we check for null (absent key) rather than emptiness.
Optional<LogSafety> cached = cache.get(typeName);
if (cached != null) {
return cached;
}
if (!inProgress.add(typeName)) {
// Given recursive evaluation, we return the least restrictive type: SAFE.
// Mark that this subtree's result depends on cycle-breaking.
encounteredCycle = true;
return OPTIONAL_OF_SAFE;
}

// Save and reset cycle state so we can detect cycles within this type's subtree only.
boolean previousCycleState = encounteredCycle;
encounteredCycle = false;

Optional<LogSafety> result = task.get();

boolean subtreeHadCycle = encounteredCycle;
// Propagate cycle detection upward: if this subtree had a cycle, callers should know.
encounteredCycle = previousCycleState || subtreeHadCycle;
Comment on lines +201 to +209
Copy link
Contributor

@aldexis aldexis Feb 24, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's slightly confusing (to me) that we store the cyclic state in a field, unset it before calling the method, then reset it afterwards.

How would you feel about returning a pair from this method e.g. record(Optional<LogSafety> safety, boolean cycleDetected) and then use that instead here? (and maybe make a second method that extracts the safety result for usage in the rest of the code)

Edit: We'd want to check the memory impact of doing these allocations though

Copy link
Author

@xanderbailey xanderbailey Feb 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I wanted to keep this lean. Want me to try that way?

Copy link
Contributor

@aldexis aldexis Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you have a bit of time to try this out, I'd be curious about the memory impact yeah (especially around the total allocations). If it's low enough, I feel like this would make the cognitive overhead of this piece of code a bit lower, and thus more maintainable.

If it's not, I think this should be fine


if (!inProgress.remove(typeName)) {
throw new IllegalStateException(
"Failed to remove " + typeName + " from in-progress, something is very wrong!");
}

// Only cache results where no cycle was encountered in the subtree.
// When a cycle is broken with the SAFE heuristic, the result depends on which type
// was the entry point, so caching it would produce incorrect results if the same
// type is later evaluated from a different starting point.
if (!subtreeHadCycle) {
cache.put(typeName, result);
}
return result;
}

Expand Down
113 changes: 113 additions & 0 deletions conjure-java-core/src/test/java/com/palantir/conjure/java/types/SafetyEvaluatorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,119 @@ void testEmptyEnum() {
.hasValue(LogSafety.SAFE);
}

@Test
void testCyclicTypes_bothSafe() {
// Foo has a field referencing Bar, Bar has a field referencing Foo.
// Both fields are marked safe, so both types should evaluate as safe.
TypeDefinition foo = TypeDefinition.object(ObjectDefinition.builder()
.typeName(FOO)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("bar"))
.type(Type.reference(BAR))
.build())
.build());
TypeDefinition bar = TypeDefinition.object(ObjectDefinition.builder()
.typeName(BAR)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("foo"))
.type(Type.reference(FOO))
.build())
.build());
Comment on lines +371 to +386
Copy link
Contributor

@aldexis aldexis Feb 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fields are not marked as safe. Shouldn't we be returning unknown/empty in this case? 🤔 (I don't think your change changes this, but this seems weird to me 🤔 )

ConjureDefinition conjureDef =
ConjureDefinition.builder().version(1).types(foo).types(bar).build();
SafetyEvaluator evaluator = new SafetyEvaluator(conjureDef);
// Evaluation order should not matter for the result
assertThat(evaluator.evaluate(foo)).hasValue(LogSafety.SAFE);
assertThat(evaluator.evaluate(bar)).hasValue(LogSafety.SAFE);
Comment on lines +389 to +392
Copy link
Contributor

@aldexis aldexis Feb 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're not checking the evaluation order here. Since you're already doing it in a dedicated test, might as well remove the comment?

Suggested change
SafetyEvaluator evaluator = new SafetyEvaluator(conjureDef);
// Evaluation order should not matter for the result
assertThat(evaluator.evaluate(foo)).hasValue(LogSafety.SAFE);
assertThat(evaluator.evaluate(bar)).hasValue(LogSafety.SAFE);
SafetyEvaluator evaluator = new SafetyEvaluator(conjureDef);
assertThat(evaluator.evaluate(foo)).hasValue(LogSafety.SAFE);
assertThat(evaluator.evaluate(bar)).hasValue(LogSafety.SAFE);

}

@Test
void testCyclicTypes_withUnsafeField() {
// Foo references Bar, Bar references Foo and also has an unsafe string field.
// The unsafe field should propagate through the cycle.
TypeDefinition foo = TypeDefinition.object(ObjectDefinition.builder()
.typeName(FOO)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("bar"))
.type(Type.reference(BAR))
.build())
.build());
TypeDefinition bar = TypeDefinition.object(ObjectDefinition.builder()
.typeName(BAR)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("foo"))
.type(Type.reference(FOO))
.build())
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("unsafeField"))
.type(Type.primitive(PrimitiveType.STRING))
.safety(LogSafety.UNSAFE)
.build())
.build());
ConjureDefinition conjureDef =
ConjureDefinition.builder().version(1).types(foo).types(bar).build();
SafetyEvaluator evaluator = new SafetyEvaluator(conjureDef);
assertThat(evaluator.evaluate(bar)).hasValue(LogSafety.UNSAFE);
assertThat(evaluator.evaluate(foo)).hasValue(LogSafety.UNSAFE);
}

@Test
void testCyclicTypes_evaluationOrderIndependent() {
// Verifies that the cache does not cause different results depending on which
// type in a cycle is evaluated first.
TypeDefinition foo = TypeDefinition.object(ObjectDefinition.builder()
.typeName(FOO)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("bar"))
.type(Type.reference(BAR))
.build())
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("unsafeField"))
.type(Type.primitive(PrimitiveType.STRING))
.safety(LogSafety.UNSAFE)
.build())
.build());
TypeDefinition bar = TypeDefinition.object(ObjectDefinition.builder()
.typeName(BAR)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("foo"))
.type(Type.reference(FOO))
.build())
.build());
ConjureDefinition conjureDef =
ConjureDefinition.builder().version(1).types(foo).types(bar).build();

// Evaluate foo first, then bar
SafetyEvaluator evaluator1 = new SafetyEvaluator(conjureDef);
Optional<LogSafety> fooFirst = evaluator1.evaluate(foo);
Optional<LogSafety> barAfterFoo = evaluator1.evaluate(bar);

// Evaluate bar first, then foo
SafetyEvaluator evaluator2 = new SafetyEvaluator(conjureDef);
Optional<LogSafety> barFirst = evaluator2.evaluate(bar);
Optional<LogSafety> fooAfterBar = evaluator2.evaluate(foo);

// Results should be identical regardless of evaluation order
assertThat(fooFirst).isEqualTo(fooAfterBar);
assertThat(barAfterFoo).isEqualTo(barFirst);
}

@Test
void testSelfReferentialType() {
// A type that references itself (e.g. a linked list node)
TypeDefinition node = TypeDefinition.object(ObjectDefinition.builder()
.typeName(FOO)
.fields(FieldDefinition.builder()
.fieldName(FieldName.of("next"))
.type(Type.reference(FOO))
.build())
.build());
ConjureDefinition conjureDef =
ConjureDefinition.builder().version(1).types(node).build();
SafetyEvaluator evaluator = new SafetyEvaluator(conjureDef);
assertThat(evaluator.evaluate(node)).hasValue(LogSafety.SAFE);
}

private static Stream<Arguments> getTypes(Type externalReference) {
TypeDefinition objectType = TypeDefinition.object(ObjectDefinition.builder()
.typeName(FOO)
Expand Down