diff --git a/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java b/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java index dcaebed..d4ddec2 100644 --- a/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java +++ b/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java @@ -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) diff --git a/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java b/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java index 18bd0a9..ed60d96 100644 --- a/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java +++ b/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java @@ -51,6 +51,8 @@ public interface EmbeddingChunkRepository extends JpaRepository 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(); + + 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( diff --git a/src/main/java/ai/decompile/ai/service/EmbeddingService.java b/src/main/java/ai/decompile/ai/service/EmbeddingService.java index 94650fb..11dd04c 100644 --- a/src/main/java/ai/decompile/ai/service/EmbeddingService.java +++ b/src/main/java/ai/decompile/ai/service/EmbeddingService.java @@ -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 chunks) { for (var i = 0; i < chunks.size(); i += batchSize) { var batch = chunks.subList(i, Math.min(i + batchSize, chunks.size())); diff --git a/src/main/java/ai/decompile/analysis/model/repository/MutationSuggestionRepository.java b/src/main/java/ai/decompile/analysis/model/repository/MutationSuggestionRepository.java index 793fc57..ad9132f 100644 --- a/src/main/java/ai/decompile/analysis/model/repository/MutationSuggestionRepository.java +++ b/src/main/java/ai/decompile/analysis/model/repository/MutationSuggestionRepository.java @@ -39,4 +39,38 @@ public interface MutationSuggestionRepository extends JpaRepository 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 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 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 findApprovedStructRenamesByAnalysisId( + @Param("analysisId") UUID analysisId); } diff --git a/src/main/java/ai/decompile/analysis/service/MutationService.java b/src/main/java/ai/decompile/analysis/service/MutationService.java index 90c7d78..9ccf4ab 100644 --- a/src/main/java/ai/decompile/analysis/service/MutationService.java +++ b/src/main/java/ai/decompile/analysis/service/MutationService.java @@ -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 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(); + + 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) diff --git a/src/main/java/ai/decompile/analysis/service/handler/FunctionTargetHandler.java b/src/main/java/ai/decompile/analysis/service/handler/FunctionTargetHandler.java index 4deda34..2ace975 100644 --- a/src/main/java/ai/decompile/analysis/service/handler/FunctionTargetHandler.java +++ b/src/main/java/ai/decompile/analysis/service/handler/FunctionTargetHandler.java @@ -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>(); 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(); if (Objects.nonNull(currentFlags) && !currentFlags.isBlank()) { diff --git a/src/main/java/ai/decompile/job/service/EmbeddingEventSubscriber.java b/src/main/java/ai/decompile/job/service/EmbeddingEventSubscriber.java index 27476db..4b38638 100644 --- a/src/main/java/ai/decompile/job/service/EmbeddingEventSubscriber.java +++ b/src/main/java/ai/decompile/job/service/EmbeddingEventSubscriber.java @@ -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; } } diff --git a/src/main/resources/db/migration/V025__alter_embedding_dimension.sql b/src/main/resources/db/migration/V025__alter_embedding_dimension.sql new file mode 100644 index 0000000..022e88b --- /dev/null +++ b/src/main/resources/db/migration/V025__alter_embedding_dimension.sql @@ -0,0 +1,2 @@ +ALTER TABLE embedding_chunk + ALTER COLUMN embedding TYPE vector(768);