[lldb-dap] Fix source references

lldb/test/API/tools/lldb-dap/breakpoint-assembly/TestDAP_breakpointAssembly.py

@@ -67,19 +67,19 @@ def test_break_on_invalid_source_reference(self):
             "Invalid sourceReference.",
         )
 
-        # Verify that setting a breakpoint on a source reference without a symbol also fails
+        # Verify that setting a breakpoint on a source reference that is not created fails
         response = self.dap_server.request_setBreakpoints(
-            Source(source_reference=0), [1]
+            Source(source_reference=200), [1]
         )
         self.assertIsNotNone(response)
         breakpoints = response["body"]["breakpoints"]
         self.assertEqual(len(breakpoints), 1)
-        breakpoint = breakpoints[0]
+        break_point = breakpoints[0]
         self.assertFalse(
-            breakpoint["verified"], "Expected breakpoint to not be verified"
+            break_point["verified"], "Expected breakpoint to not be verified"
         )
-        self.assertIn("message", breakpoint, "Expected message to be present")
+        self.assertIn("message", break_point, "Expected message to be present")
         self.assertEqual(
-            breakpoint["message"],
-            "Breakpoints in assembly without a valid symbol are not supported yet.",
+            break_point["message"],
+            "Invalid sourceReference.",
         )

lldb/tools/lldb-dap/Breakpoint.cpp

@@ -64,8 +64,8 @@ protocol::Breakpoint Breakpoint::ToProtocolBreakpoint() {
         "0x" + llvm::utohexstr(bp_addr.GetLoadAddress(m_bp.GetTarget()));
     breakpoint.instructionReference = formatted_addr;
 
-    auto source = CreateSource(bp_addr, m_dap.target);
-    if (!IsAssemblySource(source)) {
+    std::optional<protocol::Source> source = m_dap.ResolveSource(bp_addr);
+    if (source && !IsAssemblySource(*source)) {
       auto line_entry = bp_addr.GetLineEntry();
       const auto line = line_entry.GetLine();
       if (line != LLDB_INVALID_LINE_NUMBER)

lldb/tools/lldb-dap/DAP.cpp

@@ -17,6 +17,7 @@
 #include "Protocol/ProtocolBase.h"
 #include "Protocol/ProtocolRequests.h"
 #include "Protocol/ProtocolTypes.h"
+#include "ProtocolUtils.h"
 #include "Transport.h"
 #include "lldb/API/SBBreakpoint.h"
 #include "lldb/API/SBCommandInterpreter.h"
@@ -510,6 +511,27 @@ DAP::SendFormattedOutput(OutputType o, const char *format, ...) {
       o, llvm::StringRef(buffer, std::min<int>(actual_length, sizeof(buffer))));
 }
 
+int32_t DAP::CreateSourceReference(lldb::addr_t address) {
+  std::lock_guard<std::mutex> guard(m_source_references_mutex);
+  auto iter = llvm::find(m_source_references, address);
+  if (iter != m_source_references.end())
+    return std::distance(m_source_references.begin(), iter) + 1;
+
+  m_source_references.emplace_back(address);
+  return static_cast<int32_t>(m_source_references.size());
+}
+
+std::optional<lldb::addr_t> DAP::GetSourceReferenceAddress(int32_t reference) {
+  std::lock_guard<std::mutex> guard(m_source_references_mutex);
+  if (reference <= LLDB_DAP_INVALID_SRC_REF)
+    return std::nullopt;
+
+  if (static_cast<size_t>(reference) > m_source_references.size())
+    return std::nullopt;
+
+  return m_source_references[reference - 1];
+}
+
 ExceptionBreakpoint *DAP::GetExceptionBPFromStopReason(lldb::SBThread &thread) {
   const auto num = thread.GetStopReasonDataCount();
   // Check to see if have hit an exception breakpoint and change the
@@ -615,6 +637,55 @@ ReplMode DAP::DetectReplMode(lldb::SBFrame frame, std::string &expression,
   llvm_unreachable("enum cases exhausted.");
 }
 
+std::optional<protocol::Source> DAP::ResolveSource(lldb::SBAddress address) {
+  if (DisplayAssemblySource(debugger, address))
+    return ResolveAssemblySource(address);
+
+  lldb::SBLineEntry line_entry = GetLineEntryForAddress(target, address);
+  if (!line_entry.IsValid())
+    return std::nullopt;
+
+  return CreateSource(line_entry.GetFileSpec());
+}
+
+std::optional<protocol::Source>
+DAP::ResolveAssemblySource(lldb::SBAddress address) {
+  lldb::SBSymbol symbol = address.GetSymbol();
+  lldb::addr_t load_addr = LLDB_INVALID_ADDRESS;
+  std::string name;
+  if (symbol.IsValid()) {
+    load_addr = symbol.GetStartAddress().GetLoadAddress(target);
+    name = symbol.GetName();
+  } else {
+    load_addr = address.GetLoadAddress(target);
+    name = GetLoadAddressString(load_addr);
+  }
+
+  if (load_addr == LLDB_INVALID_ADDRESS)
+    return std::nullopt;
+
+  protocol::Source source;
+  source.sourceReference = CreateSourceReference(load_addr);
+  lldb::SBModule module = address.GetModule();
+  if (module.IsValid()) {
+    lldb::SBFileSpec file_spec = module.GetFileSpec();
+    if (file_spec.IsValid()) {
+      std::string path = GetSBFileSpecPath(file_spec);
+      if (!path.empty())
+        source.path = path + '`' + name;
+    }
+  }
+
+  source.name = std::move(name);
+
+  // Mark the source as deemphasized since users will only be able to view
+  // assembly for these frames.
+  source.presentationHint =
+      protocol::Source::eSourcePresentationHintDeemphasize;
+
+  return source;
+}
+
 bool DAP::RunLLDBCommands(llvm::StringRef prefix,
                           llvm::ArrayRef<std::string> commands) {
   bool required_command_failed = false;

lldb/tools/lldb-dap/DAP.h

@@ -219,7 +219,9 @@ struct DAP {
   void __attribute__((format(printf, 3, 4)))
   SendFormattedOutput(OutputType o, const char *format, ...);
 
-  static int64_t GetNextSourceReference();
+  int32_t CreateSourceReference(lldb::addr_t address);
+
+  std::optional<lldb::addr_t> GetSourceReferenceAddress(int32_t reference);
 
   ExceptionBreakpoint *GetExceptionBPFromStopReason(lldb::SBThread &thread);
 
@@ -252,6 +254,29 @@ struct DAP {
   ReplMode DetectReplMode(lldb::SBFrame frame, std::string &expression,
                           bool partial_expression);
 
+  /// Create a "Source" JSON object as described in the debug adapter
+  /// definition.
+  ///
+  /// \param[in] address
+  ///     The address to use when populating out the "Source" object.
+  ///
+  /// \return
+  ///     An optional "Source" JSON object that follows the formal JSON
+  ///     definition outlined by Microsoft.
+  std::optional<protocol::Source> ResolveSource(lldb::SBAddress address);
+
+  /// Create a "Source" JSON object as described in the debug adapter
+  /// definition.
+  ///
+  /// \param[in] address
+  ///     The address to use when populating out the "Source" object.
+  ///
+  /// \return
+  ///     An optional "Source" JSON object that follows the formal JSON
+  ///     definition outlined by Microsoft.
+  std::optional<protocol::Source>
+  ResolveAssemblySource(lldb::SBAddress address);
+
   /// \return
   ///   \b false if a fatal error was found while executing these commands,
   ///   according to the rules of \a LLDBUtils::RunLLDBCommands.
@@ -406,6 +431,10 @@ struct DAP {
   std::thread progress_event_thread;
   /// @}
 
+  /// List of addresses mapped by sourceReference.
+  std::vector<lldb::addr_t> m_source_references;
+  std::mutex m_source_references_mutex;
+
   /// Queue for all incoming messages.
   std::deque<protocol::Message> m_queue;
   std::mutex m_queue_mutex;

lldb/tools/lldb-dap/Handler/DisassembleRequestHandler.cpp

@@ -85,7 +85,8 @@ static lldb::SBAddress GetDisassembleStartAddress(lldb::SBTarget target,
 }
 
 static DisassembledInstruction ConvertSBInstructionToDisassembledInstruction(
-    lldb::SBTarget &target, lldb::SBInstruction &inst, bool resolve_symbols) {
+    DAP &dap, lldb::SBInstruction &inst, bool resolve_symbols) {
+  lldb::SBTarget target = dap.target;
   if (!inst.IsValid())
     return GetInvalidInstruction();
 
@@ -138,14 +139,14 @@ static DisassembledInstruction ConvertSBInstructionToDisassembledInstruction(
     si << " ; " << c;
   }
 
-  protocol::Source source = CreateSource(addr, target);
+  std::optional<protocol::Source> source = dap.ResolveSource(addr);
   lldb::SBLineEntry line_entry = GetLineEntryForAddress(target, addr);
 
   // If the line number is 0 then the entry represents a compiler generated
   // location.
-  if (!IsAssemblySource(source) && line_entry.GetStartAddress() == addr &&
-      line_entry.IsValid() && line_entry.GetFileSpec().IsValid() &&
-      line_entry.GetLine() != 0) {
+  if (source && !IsAssemblySource(*source) &&
+      line_entry.GetStartAddress() == addr && line_entry.IsValid() &&
+      line_entry.GetFileSpec().IsValid() && line_entry.GetLine() != 0) {
 
     disassembled_inst.location = std::move(source);
     const auto line = line_entry.GetLine();
@@ -221,7 +222,7 @@ DisassembleRequestHandler::Run(const DisassembleArguments &args) const {
       original_address_index = i;
 
     instructions.push_back(ConvertSBInstructionToDisassembledInstruction(
-        dap.target, inst, resolve_symbols));
+        dap, inst, resolve_symbols));
   }
 
   // Check if we miss instructions at the beginning.

lldb/tools/lldb-dap/Handler/LocationsRequestHandler.cpp

@@ -137,7 +137,16 @@ void LocationsRequestHandler::operator()(
       return;
     }
 
-    body.try_emplace("source", CreateSource(line_entry.GetFileSpec()));
+    const std::optional<protocol::Source> source =
+        CreateSource(line_entry.GetFileSpec());
+    if (!source) {
+      response["success"] = false;
+      response["message"] = "Failed to resolve file path for location";
+      dap.SendJSON(llvm::json::Value(std::move(response)));
+      return;
+    }
+
+    body.try_emplace("source", *source);
     if (int line = line_entry.GetLine())
       body.try_emplace("line", line);
     if (int column = line_entry.GetColumn())
@@ -152,7 +161,16 @@ void LocationsRequestHandler::operator()(
       return;
     }
 
-    body.try_emplace("source", CreateSource(decl.GetFileSpec()));
+    const std::optional<protocol::Source> source =
+        CreateSource(decl.GetFileSpec());
+    if (!source) {
+      response["success"] = false;
+      response["message"] = "Failed to resolve file path for location";
+      dap.SendJSON(llvm::json::Value(std::move(response)));
+      return;
+    }
+
+    body.try_emplace("source", *source);
     if (int line = decl.GetLine())
       body.try_emplace("line", line);
     if (int column = decl.GetColumn())

lldb/tools/lldb-dap/Handler/SourceRequestHandler.cpp

@@ -29,34 +29,42 @@ namespace lldb_dap {
 /// the source code for a given source reference.
 llvm::Expected<protocol::SourceResponseBody>
 SourceRequestHandler::Run(const protocol::SourceArguments &args) const {
-  const auto source =
+
+  uint32_t source_ref =
       args.source->sourceReference.value_or(args.sourceReference);
+  const std::optional<lldb::addr_t> source_addr_opt =
+      dap.GetSourceReferenceAddress(source_ref);
 
-  if (!source)
+  if (!source_addr_opt)
     return llvm::make_error<DAPError>(
-        "invalid arguments, expected source.sourceReference to be set");
+        llvm::formatv("Unknown source reference {}", source_ref));
 
-  lldb::SBAddress address(source, dap.target);
+  lldb::SBAddress address(*source_addr_opt, dap.target);
   if (!address.IsValid())
     return llvm::make_error<DAPError>("source not found");
 
   lldb::SBSymbol symbol = address.GetSymbol();
-
-  lldb::SBStream stream;
-  lldb::SBExecutionContext exe_ctx(dap.target);
+  lldb::SBInstructionList insts;
 
   if (symbol.IsValid()) {
-    lldb::SBInstructionList insts = symbol.GetInstructions(dap.target);
-    insts.GetDescription(stream, exe_ctx);
+    insts = symbol.GetInstructions(dap.target);
   } else {
     // No valid symbol, just return the disassembly.
-    lldb::SBInstructionList insts = dap.target.ReadInstructions(
+    insts = dap.target.ReadInstructions(
         address, dap.k_number_of_assembly_lines_for_nodebug);
-    insts.GetDescription(stream, exe_ctx);
   }
 
+  if (!insts || insts.GetSize() == 0)
+    return llvm::make_error<DAPError>(
+        llvm::formatv("no instruction source for address {}",
+                      address.GetLoadAddress(dap.target)));
+
+  lldb::SBStream stream;
+  lldb::SBExecutionContext exe_ctx(dap.target);
+  insts.GetDescription(stream, exe_ctx);
   return protocol::SourceResponseBody{/*content=*/stream.GetData(),
-                                      /*mimeType=*/"text/x-lldb.disassembly"};
+                                      /*mimeType=*/
+                                      "text/x-lldb.disassembly"};
 }
 
 } // namespace lldb_dap

lldb/tools/lldb-dap/Handler/StackTraceRequestHandler.cpp

@@ -69,7 +69,7 @@ static bool FillStackFrames(DAP &dap, lldb::SBThread &thread,
       break;
     }
 
-    stack_frames.emplace_back(CreateStackFrame(frame, frame_format));
+    stack_frames.emplace_back(CreateStackFrame(dap, frame, frame_format));
   }
 
   if (include_all && reached_end_of_stack) {

lldb/tools/lldb-dap/JSONUtils.cpp

@@ -552,7 +552,7 @@ CreateExceptionBreakpointFilter(const ExceptionBreakpoint &bp) {
 //   },
 //   "required": [ "id", "name", "line", "column" ]
 // }
-llvm::json::Value CreateStackFrame(lldb::SBFrame &frame,
+llvm::json::Value CreateStackFrame(DAP &dap, lldb::SBFrame &frame,
                                    lldb::SBFormat &format) {
   llvm::json::Object object;
   int64_t frame_id = MakeDAPFrameID(frame);
@@ -582,8 +582,10 @@ llvm::json::Value CreateStackFrame(lldb::SBFrame &frame,
   EmplaceSafeString(object, "name", frame_name);
 
   auto target = frame.GetThread().GetProcess().GetTarget();
-  auto source = CreateSource(frame.GetPCAddress(), target);
-  if (!IsAssemblySource(source)) {
+  std::optional<protocol::Source> source =
+      dap.ResolveSource(frame.GetPCAddress());
+
+  if (source && !IsAssemblySource(*source)) {
     // This is a normal source with a valid line entry.
     auto line_entry = frame.GetLineEntry();
     object.try_emplace("line", line_entry.GetLine());
@@ -607,7 +609,8 @@ llvm::json::Value CreateStackFrame(lldb::SBFrame &frame,
     object.try_emplace("column", 1);
   }
 
-  object.try_emplace("source", std::move(source));
+  if (source)
+    object.try_emplace("source", std::move(source).value());
 
   const auto pc = frame.GetPC();
   if (pc != LLDB_INVALID_ADDRESS) {

lldb/tools/lldb-dap/JSONUtils.h

@@ -246,6 +246,9 @@ CreateExceptionBreakpointFilter(const ExceptionBreakpoint &bp);
 ///   "line" - the source file line number as an integer
 ///   "column" - the source file column number as an integer
 ///
+/// \param[in] dap
+///     The DAP session associated with the stopped thread.
+///
 /// \param[in] frame
 ///     The LLDB stack frame to use when populating out the "StackFrame"
 ///     object.
@@ -257,7 +260,7 @@ CreateExceptionBreakpointFilter(const ExceptionBreakpoint &bp);
 /// \return
 ///     A "StackFrame" JSON object with that follows the formal JSON
 ///     definition outlined by Microsoft.
-llvm::json::Value CreateStackFrame(lldb::SBFrame &frame,
+llvm::json::Value CreateStackFrame(DAP &dap, lldb::SBFrame &frame,
                                    lldb::SBFormat &format);
 
 /// Create a "StackFrame" label object for a LLDB thread.

lldb/tools/lldb-dap/Protocol/ProtocolRequests.h

@@ -468,7 +468,7 @@ struct SourceArguments {
   /// The reference to the source. This is the same as `source.sourceReference`.
   /// This is provided for backward compatibility since old clients do not
   /// understand the `source` attribute.
-  int64_t sourceReference;
+  int64_t sourceReference = LLDB_DAP_INVALID_SRC_REF;
 };
 bool fromJSON(const llvm::json::Value &, SourceArguments &, llvm::json::Path);
 

lldb/tools/lldb-dap/Protocol/ProtocolTypes.cpp

@@ -65,7 +65,7 @@ llvm::json::Value toJSON(const Source &S) {
     result.insert({"name", *S.name});
   if (S.path)
     result.insert({"path", *S.path});
-  if (S.sourceReference)
+  if (S.sourceReference && (*S.sourceReference > LLDB_DAP_INVALID_SRC_REF))
     result.insert({"sourceReference", *S.sourceReference});
   if (S.presentationHint)
     result.insert({"presentationHint", *S.presentationHint});