From 54f34d72ad92a4f8c7a4d9772536eb61660b1af8 Mon Sep 17 00:00:00 2001
From: Gavin Nishizawa <gavin@terrastruct.com>
Date: Wed, 11 Jan 2023 20:21:47 -0800
Subject: [PATCH] extend previous segment to target if last segment is very
 short

---
 d2layouts/d2dagrelayout/layout.go | 43 +++++++++++++++++++++++++++++++
 lib/geo/point.go                  |  6 ++---
 lib/geo/point_test.go             |  2 +-
 lib/geo/segment.go                |  8 ++++++
 4 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/d2layouts/d2dagrelayout/layout.go b/d2layouts/d2dagrelayout/layout.go
index 34b87fcd6..7ac2a44e4 100644
--- a/d2layouts/d2dagrelayout/layout.go
+++ b/d2layouts/d2dagrelayout/layout.go
@@ -30,6 +30,8 @@ var setupJS string
 //go:embed dagre.js
 var dagreJS string
 
+const MIN_SEGMENT_LEN = 10
+
 type ConfigurableOpts struct {
 	NodeSep int `json:"nodesep"`
 	EdgeSep int `json:"edgesep"`
@@ -247,6 +249,47 @@ func Layout(ctx context.Context, g *d2graph.Graph, opts *ConfigurableOpts) (err
 			}
 		}
 
+		// arrowheads can appear broken if segments are very short from dagre routing a point just outside the shape
+		// to fix this, we try extending the previous segment into the shape instead of having a very short segment
+		if !start.Equals(points[0]) && startIndex+2 < len(points) {
+			newStartingSegment := *geo.NewSegment(start, points[startIndex+1])
+			if newStartingSegment.Length() < MIN_SEGMENT_LEN {
+				// we don't want a very short segment right next to the source because it will mess up the arrowhead
+				// instead we want to extend the next segment into the shape border if possible
+				nextStart := points[startIndex+1]
+				nextEnd := points[startIndex+2]
+
+				// Note: in other direction to extend towards source
+				nextSegment := *geo.NewSegment(nextStart, nextEnd)
+				v := nextSegment.ToVector()
+				extendedStart := nextEnd.ToVector().Add(v.AddLength(MIN_SEGMENT_LEN)).ToPoint()
+				extended := *geo.NewSegment(nextEnd, extendedStart)
+
+				if intersections := edge.Src.Box.Intersections(extended); len(intersections) > 0 {
+					start = intersections[0]
+					startIndex += 1
+				}
+			}
+		}
+		if !end.Equals(points[len(points)-1]) && endIndex-2 >= 0 {
+			newEndingSegment := *geo.NewSegment(end, points[endIndex-1])
+			if newEndingSegment.Length() < MIN_SEGMENT_LEN {
+				// extend the prev segment into the shape border if possible
+				prevStart := points[endIndex-2]
+				prevEnd := points[endIndex-1]
+
+				prevSegment := *geo.NewSegment(prevStart, prevEnd)
+				v := prevSegment.ToVector()
+				extendedEnd := prevStart.ToVector().Add(v.AddLength(MIN_SEGMENT_LEN)).ToPoint()
+				extended := *geo.NewSegment(prevStart, extendedEnd)
+
+				if intersections := edge.Dst.Box.Intersections(extended); len(intersections) > 0 {
+					end = intersections[0]
+					endIndex -= 1
+				}
+			}
+		}
+
 		srcShape := shape.NewShape(d2target.DSL_SHAPE_TO_SHAPE_TYPE[strings.ToLower(edge.Src.Attributes.Shape.Value)], edge.Src.Box)
 		dstShape := shape.NewShape(d2target.DSL_SHAPE_TO_SHAPE_TYPE[strings.ToLower(edge.Dst.Attributes.Shape.Value)], edge.Dst.Box)
 
diff --git a/lib/geo/point.go b/lib/geo/point.go
index 13275db79..0acc5021b 100644
--- a/lib/geo/point.go
+++ b/lib/geo/point.go
@@ -187,12 +187,12 @@ func (p *Point) DistanceToLine(p1, p2 *Point) float64 {
 
 // Moves the given point by Vector
 func (start *Point) AddVector(v Vector) *Point {
-	return start.toVector().Add(v).ToPoint()
+	return start.ToVector().Add(v).ToPoint()
 }
 
 // Creates a Vector of the size between start and endpoint, pointing to endpoint
 func (start *Point) VectorTo(endpoint *Point) Vector {
-	return endpoint.toVector().Minus(start.toVector())
+	return endpoint.ToVector().Minus(start.ToVector())
 }
 
 func (p *Point) FormattedCoordinates() string {
@@ -205,7 +205,7 @@ func (q *Point) OnSegment(p, r *Point) bool {
 }
 
 // Creates a Vector pointing to point
-func (endpoint *Point) toVector() Vector {
+func (endpoint *Point) ToVector() Vector {
 	return []float64{endpoint.X, endpoint.Y}
 }
 
diff --git a/lib/geo/point_test.go b/lib/geo/point_test.go
index 259975897..343f8bb30 100644
--- a/lib/geo/point_test.go
+++ b/lib/geo/point_test.go
@@ -29,7 +29,7 @@ func TestAddVector(t *testing.T) {
 
 func TestToVector(t *testing.T) {
 	p := &Point{3.5, 6.7}
-	v := p.toVector()
+	v := p.ToVector()
 
 	if v[0] != p.X || v[1] != p.Y {
 		t.Fatalf("Expected Vector (%v) coordinates to match the point (%v)", p, v)
diff --git a/lib/geo/segment.go b/lib/geo/segment.go
index 473cb12e8..df6f64a5c 100644
--- a/lib/geo/segment.go
+++ b/lib/geo/segment.go
@@ -114,3 +114,11 @@ func (segment *Segment) GetBounds(segments []*Segment, buffer float64) (float64,
 	}
 	return floor, ceil
 }
+
+func (segment Segment) Length() float64 {
+	return EuclideanDistance(segment.Start.X, segment.Start.Y, segment.End.X, segment.End.Y)
+}
+
+func (segment Segment) ToVector() Vector {
+	return NewVector(segment.End.X-segment.Start.X, segment.End.Y-segment.Start.Y)
+}