Decompile.AI/backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: rename maps for labels/structs/globals, reindex triggers, comment validation

Browse Source
This commit is contained in:
Rodrigo Verdiani 2026-07-04 15:38:37 -03:00
parent 75275f992d
commit 0c54602fcf
9 changed files with 160 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java
View File

@ -39,7 +39,7 @@ public class EmbeddingChunk {
private String content;
@JdbcTypeCode(SqlTypes.VECTOR)
@Column(name = "embedding", columnDefinition = "vector(4096)")
@Column(name = "embedding", columnDefinition = "vector(768)")
private float[] embedding;
@JdbcTypeCode(SqlTypes.JSON)

2
src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java
View File

@ -51,6 +51,8 @@ public interface EmbeddingChunkRepository extends JpaRepository<EmbeddingChunk,
void deleteByBinaryId(UUID binaryId);
void deleteByBinaryIdAndChunkType(UUID binaryId, String chunkType);
@Modifying
@Transactional
@Query("DELETE FROM EmbeddingChunk e WHERE e.sourceId = :sourceId AND e.sourceType = :sourceType")

30
src/main/java/ai/decompile/ai/service/EmbeddingChunkBuilder.java
View File

@ -1,7 +1,6 @@
package ai.decompile.ai.service;
import ai.decompile.ai.model.entity.EmbeddingChunk;
import ai.decompile.analysis.model.entity.MutationSuggestion;
import ai.decompile.analysis.model.entity.StaticAnalysis;
import ai.decompile.analysis.model.entity.StaticDataLabel;
import ai.decompile.analysis.model.entity.StaticEnum;
@ -87,10 +86,31 @@ public class EmbeddingChunkBuilder {
}
var analysisId = functions.getFirst().getAnalysis().getId();
return mutationRepository.findApprovedFunctionRenamesByAnalysisId(analysisId).stream()
.collect(
java.util.stream.Collectors.toMap(
MutationSuggestion::getOldValue, MutationSuggestion::getNewValue, (_, b) -> b));
var functionRenames = mutationRepository.findApprovedFunctionRenamesByAnalysisId(analysisId);
var labelRenames = mutationRepository.findApprovedLabelRenamesByAnalysisId(analysisId);
var globalVarRenames =
mutationRepository.findApprovedGlobalVarRenamesByAnalysisId(analysisId);
var structRenames = mutationRepository.findApprovedStructRenamesByAnalysisId(analysisId);
var map = new HashMap<String, String>();
for (var m : labelRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : globalVarRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : structRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : functionRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
return map;
}
public EmbeddingChunk buildMetadataChunk(

26
src/main/java/ai/decompile/ai/service/EmbeddingService.java
View File

@ -133,6 +133,32 @@ public class EmbeddingService {
binaryId);
}
@Transactional
public void reindexAllFunctions(UUID binaryId) {
var analysis = getLatestAnalysis(binaryId);
if (Objects.isNull(analysis)) {
log.warn("reindexAllFunctions: no analysis found for binary {}, skipping", binaryId);
return;
}
var functions = analysis.getFunctions();
if (functions.isEmpty()) {
log.debug("reindexAllFunctions: no functions for binary {}, nothing to do", binaryId);
return;
}
var binary = binaryService.getBinary(binaryId);
embeddingChunkRepository.deleteByBinaryIdAndChunkType(binaryId, "FUNCTION");
embeddingChunkRepository.flush();
var chunks = chunkBuilder.buildFunctionChunks(functions, binary);
embedAndPersistChunks(chunks);
log.info(
"Re-indexed all {} function chunks for binary {}", functions.size(), binaryId);
}
private void embedAndPersistChunks(List<EmbeddingChunk> chunks) {
for (var i = 0; i < chunks.size(); i += batchSize) {
var batch = chunks.subList(i, Math.min(i + batchSize, chunks.size()));

34
src/main/java/ai/decompile/analysis/model/repository/MutationSuggestionRepository.java
View File

@ -39,4 +39,38 @@ public interface MutationSuggestionRepository extends JpaRepository<MutationSugg
""")
List<MutationSuggestion> findApprovedFunctionRenamesByAnalysisId(
@Param("analysisId") UUID analysisId);
@Query(
"""
SELECT m FROM MutationSuggestion m
JOIN StaticLabel l ON l.id = m.targetId
JOIN StaticFunction f ON f.id = l.function.id
JOIN StaticAnalysis a ON a.id = f.analysis.id
WHERE m.targetType = 'LABEL' AND m.fieldName = 'name'
AND a.id = :analysisId AND m.status = 'APPROVED'
""")
List<MutationSuggestion> findApprovedLabelRenamesByAnalysisId(
@Param("analysisId") UUID analysisId);
@Query(
"""
SELECT m FROM MutationSuggestion m
JOIN StaticDataLabel l ON l.id = m.targetId
JOIN StaticAnalysis a ON a.id = l.analysis.id
WHERE m.targetType = 'GLOBAL_VAR' AND m.fieldName = 'name'
AND a.id = :analysisId AND m.status = 'APPROVED'
""")
List<MutationSuggestion> findApprovedGlobalVarRenamesByAnalysisId(
@Param("analysisId") UUID analysisId);
@Query(
"""
SELECT m FROM MutationSuggestion m
JOIN StaticStruct s ON s.id = m.targetId
JOIN StaticAnalysis a ON a.id = s.analysis.id
WHERE m.targetType = 'STRUCT' AND m.fieldName = 'name'
AND a.id = :analysisId AND m.status = 'APPROVED'
""")
List<MutationSuggestion> findApprovedStructRenamesByAnalysisId(
@Param("analysisId") UUID analysisId);
}

33
src/main/java/ai/decompile/analysis/service/MutationService.java
View File

@ -16,6 +16,7 @@ import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.stream.Collectors;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.context.ApplicationEventPublisher;
@ -175,8 +176,7 @@ public class MutationService {
var fieldName = "comments";
var oldValue = handler.getCurrentValue(functionId, fieldName);
handler.applyComment(functionId, comment, address);
var newValue = handler.getCurrentValue(functionId, fieldName);
var newValue = handler.computeUpdatedComments(oldValue, comment, address);
var mutation =
buildMutation(targetType, functionId, fieldName, oldValue, newValue, source, "PENDING");
@ -285,10 +285,31 @@ public class MutationService {
@Transactional(readOnly = true)
public Map<String, String> buildRenameMap(UUID analysisId) {
return mutationRepository.findApprovedFunctionRenamesByAnalysisId(analysisId).stream()
.collect(
Collectors.toMap(
MutationSuggestion::getOldValue, MutationSuggestion::getNewValue, (_, b) -> b));
var functionRenames = mutationRepository.findApprovedFunctionRenamesByAnalysisId(analysisId);
var labelRenames = mutationRepository.findApprovedLabelRenamesByAnalysisId(analysisId);
var globalVarRenames =
mutationRepository.findApprovedGlobalVarRenamesByAnalysisId(analysisId);
var structRenames = mutationRepository.findApprovedStructRenamesByAnalysisId(analysisId);
var map = new java.util.HashMap<String, String>();
for (var m : labelRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : globalVarRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : structRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
for (var m : functionRenames) {
map.put(m.getOldValue(), m.getNewValue());
}
return map;
}
@Transactional(readOnly = true)

22
src/main/java/ai/decompile/analysis/service/handler/FunctionTargetHandler.java
View File

@ -52,7 +52,13 @@ public class FunctionTargetHandler implements MutationTargetHandler {
switch (fieldName) {
case "name" -> func.setName(newValue);
case "signature" -> func.setSignature(newValue);
case "comments" -> func.setComments(appendComment(func.getComments(), newValue, ""));
case "comments" -> {
if ("SET".equals(operation)) {
func.setComments(newValue);
} else {
func.setComments(appendComment(func.getComments(), newValue, ""));
}
}
case "flags" -> func.setFlags(applyFlagOperation(func.getFlags(), newValue, operation));
default -> throw new IllegalArgumentException("Unknown field: " + fieldName);
}
@ -72,7 +78,13 @@ public class FunctionTargetHandler implements MutationTargetHandler {
functionRepository.save(func);
}
public String computeUpdatedComments(String existingComments, String comment, String address) {
return appendComment(existingComments, comment, address);
}
private String appendComment(String existingComments, String newCommentText, String address) {
validateNotJsonArray(newCommentText);
var list = new ArrayList<Map<String, Object>>();
if (Objects.nonNull(existingComments)
@ -104,6 +116,14 @@ public class FunctionTargetHandler implements MutationTargetHandler {
}
}
private void validateNotJsonArray(String commentText) {
var trimmed = commentText.trim();
if (trimmed.startsWith("[") && trimmed.endsWith("]")) {
throw new IllegalArgumentException(
"Comment text looks like a JSON array. Pass only the comment text, not the full comments array.");
}
}
private String applyFlagOperation(String currentFlags, String flag, String operation) {
var parts = new ArrayList<String>();
if (Objects.nonNull(currentFlags) && !currentFlags.isBlank()) {

24
src/main/java/ai/decompile/job/service/EmbeddingEventSubscriber.java
View File

@ -3,11 +3,13 @@ package ai.decompile.job.service;
import ai.decompile.ai.service.EmbeddingService;
import ai.decompile.analysis.event.MutationAppliedEvent;
import ai.decompile.analysis.event.StaticAnalysisCompletedEvent;
import ai.decompile.analysis.model.repository.StaticLabelRepository;
import ai.decompile.job.messaging.JobMessagePublisher;
import ai.decompile.job.model.entity.Job;
import ai.decompile.job.model.enums.JobStatus;
import ai.decompile.job.model.enums.JobType;
import ai.decompile.job.model.repository.JobRepository;
import java.util.Objects;
import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
@ -24,6 +26,7 @@ public class EmbeddingEventSubscriber {
private final JobRepository jobRepository;
private final JobMessagePublisher jobMessagePublisher;
private final EmbeddingService embeddingService;
private final StaticLabelRepository labelRepository;
@ApplicationModuleListener
public void on(StaticAnalysisCompletedEvent event) {
@ -48,8 +51,25 @@ public class EmbeddingEventSubscriber {
event.fieldName(),
event.targetId());
if ("FUNCTION".equals(event.targetType()) && "name".equals(event.fieldName())) {
embeddingService.reindexFunction(event.targetId(), event.binaryId());
if (!"name".equals(event.fieldName())) {
return;
}
switch (event.targetType()) {
case "FUNCTION":
embeddingService.reindexFunction(event.targetId(), event.binaryId());
break;
case "LABEL":
var label = labelRepository.findById(event.targetId()).orElse(null);
if (Objects.nonNull(label)) {
embeddingService.reindexFunction(label.getFunction().getId(), event.binaryId());
}
break;
case "GLOBAL_VAR", "STRUCT":
embeddingService.reindexAllFunctions(event.binaryId());
break;
default:
break;
}
}

2
src/main/resources/db/migration/V025__alter_embedding_dimension.sql Normal file
View File

@ -0,0 +1,2 @@
ALTER TABLE embedding_chunk
ALTER COLUMN embedding TYPE vector(768);