From 2f71d575df4390579eff774b78e38d1d459f5cd1 Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Thu, 28 Sep 2023 17:54:01 -0700
Subject: [PATCH] handle lifeline edges

---
 d2layouts/d2layouts.go                   |  9 +++++++++
 d2layouts/d2sequence/sequence_diagram.go | 20 ++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/d2layouts/d2layouts.go b/d2layouts/d2layouts.go
index adfd8eea5..d772b0774 100644
--- a/d2layouts/d2layouts.go
+++ b/d2layouts/d2layouts.go
@@ -318,6 +318,15 @@ func ExtractSubgraph(container *d2graph.Object, includeSelf bool) (nestedGraph *
 	remainingEdges := make([]*d2graph.Edge, 0, len(g.Edges))
 	for _, edge := range g.Edges {
 		srcIsNested := isNestedObject(edge.Src)
+		if d2sequence.IsLifelineEnd(edge.Dst) {
+			// special handling for lifelines since their edge.Dst is a special Object
+			if srcIsNested {
+				nestedGraph.Edges = append(nestedGraph.Edges, edge)
+			} else {
+				remainingEdges = append(remainingEdges, edge)
+			}
+			continue
+		}
 		dstIsNested := isNestedObject(edge.Dst)
 		if srcIsNested && dstIsNested {
 			nestedGraph.Edges = append(nestedGraph.Edges, edge)
diff --git a/d2layouts/d2sequence/sequence_diagram.go b/d2layouts/d2sequence/sequence_diagram.go
index ac2e8d2e2..aaaca5d2b 100644
--- a/d2layouts/d2sequence/sequence_diagram.go
+++ b/d2layouts/d2sequence/sequence_diagram.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"strconv"
 	"strings"
 
 	"oss.terrastruct.com/util-go/go2"
@@ -411,6 +412,25 @@ func (sd *sequenceDiagram) addLifelineEdges() {
 	}
 }
 
+func IsLifelineEnd(obj *d2graph.Object) bool {
+	// lifeline ends only have ID and no graph parent or box set
+	if obj.Graph != nil || obj.Parent != nil || obj.Box != nil {
+		return false
+	}
+	if !strings.Contains(obj.ID, "-lifeline-end-") {
+		return false
+	}
+	parts := strings.Split(obj.ID, "-lifeline-end-")
+	if len(parts) > 1 {
+		hash := parts[len(parts)-1]
+		actorID := strings.Join(parts[:len(parts)-1], "-lifeline-end-")
+		if strconv.Itoa(go2.StringToIntHash(actorID+"-lifeline-end")) == hash {
+			return true
+		}
+	}
+	return false
+}
+
 func (sd *sequenceDiagram) placeNotes() {
 	rankToX := make(map[int]float64)
 	for _, actor := range sd.actors {