Add arbitrary length intersections

bluesky-social · Sep 16, 2023 · c60c8d1 · c60c8d1
1 parent b889cc0
commit c60c8d1
Show file tree

Hide file tree

Showing 3 changed files with 150 additions and 68 deletions.
diff --git a/graphd/graph.go b/graphd/graph.go
@@ -198,71 +198,97 @@ func (g *Graph) GetFollowers(uid uint64) ([]uint64, error) {
 	return followers, nil
 }
 
-func (g *Graph) IntersectFollowers(uidA, uidB uint64) ([]uint64, error) {
-	followMapA, ok := g.followers.Load(uidA)
-	if !ok {
-		return nil, fmt.Errorf("uid %d not found", uidA)
+func (g *Graph) IntersectFollowers(uids []uint64) ([]uint64, error) {
+	if len(uids) == 0 {
+		return nil, fmt.Errorf("uids must have at least one element")
 	}
-	followMapA.(*FollowMap).lk.RLock()
-	defer followMapA.(*FollowMap).lk.RUnlock()
 
-	followMapB, ok := g.followers.Load(uidB)
-	if !ok {
-		return nil, fmt.Errorf("uid %d not found", uidB)
+	if len(uids) == 1 {
+		return g.GetFollowers(uids[0])
 	}
-	followMapB.(*FollowMap).lk.RLock()
-	defer followMapB.(*FollowMap).lk.RUnlock()
 
-	intersection := make([]uint64, 0)
+	followMaps := make([]*FollowMap, len(uids))
+	for i, uid := range uids {
+		followMap, ok := g.followers.Load(uid)
+		if !ok {
+			return nil, fmt.Errorf("uid %d not found", uid)
+		}
+		followMap.(*FollowMap).lk.RLock()
+		defer followMap.(*FollowMap).lk.RUnlock()
+		followMaps[i] = followMap.(*FollowMap)
+	}
 
-	// Iterate over the smaller map to speed up intersections between asymmetric sets
-	if len(followMapA.(*FollowMap).data) < len(followMapB.(*FollowMap).data) {
-		for follower := range followMapA.(*FollowMap).data {
-			if _, ok := followMapB.(*FollowMap).data[follower]; ok {
-				intersection = append(intersection, follower)
-			}
+	// Find the smallest map
+	smallest := followMaps[0]
+	for _, followMap := range followMaps {
+		if len(followMap.data) < len(smallest.data) {
+			smallest = followMap
 		}
-	} else {
-		for follower := range followMapB.(*FollowMap).data {
-			if _, ok := followMapA.(*FollowMap).data[follower]; ok {
-				intersection = append(intersection, follower)
+	}
+
+	// Remove the smallest map from the list of maps to intersect
+	followMaps = followMaps[1:]
+
+	intersection := make([]uint64, 0)
+	for follower := range smallest.data {
+		found := true
+		for _, followMap := range followMaps {
+			if _, ok := followMap.data[follower]; !ok {
+				found = false
+				break
 			}
 		}
+		if found {
+			intersection = append(intersection, follower)
+		}
 	}
 
 	return intersection, nil
 }
 
-func (g *Graph) IntersectFollowing(uidA, uidB uint64) ([]uint64, error) {
-	followMapA, ok := g.follows.Load(uidA)
-	if !ok {
-		return nil, fmt.Errorf("uid %d not found", uidA)
+func (g *Graph) IntersectFollowing(uids []uint64) ([]uint64, error) {
+	if len(uids) == 0 {
+		return nil, fmt.Errorf("uids must have at least one element")
 	}
-	followMapA.(*FollowMap).lk.RLock()
-	defer followMapA.(*FollowMap).lk.RUnlock()
 
-	followMapB, ok := g.follows.Load(uidB)
-	if !ok {
-		return nil, fmt.Errorf("uid %d not found", uidB)
+	if len(uids) == 1 {
+		return g.GetFollowing(uids[0])
 	}
-	followMapB.(*FollowMap).lk.RLock()
-	defer followMapB.(*FollowMap).lk.RUnlock()
 
-	intersection := make([]uint64, 0)
+	followMaps := make([]*FollowMap, len(uids))
+	for i, uid := range uids {
+		followMap, ok := g.follows.Load(uid)
+		if !ok {
+			return nil, fmt.Errorf("uid %d not found", uid)
+		}
+		followMap.(*FollowMap).lk.RLock()
+		defer followMap.(*FollowMap).lk.RUnlock()
+		followMaps[i] = followMap.(*FollowMap)
+	}
 
-	// Iterate over the smaller map to speed up intersections between asymmetric sets
-	if len(followMapA.(*FollowMap).data) < len(followMapB.(*FollowMap).data) {
-		for following := range followMapA.(*FollowMap).data {
-			if _, ok := followMapB.(*FollowMap).data[following]; ok {
-				intersection = append(intersection, following)
-			}
+	// Find the smallest map
+	smallest := followMaps[0]
+	for _, followMap := range followMaps {
+		if len(followMap.data) < len(smallest.data) {
+			smallest = followMap
 		}
-	} else {
-		for following := range followMapB.(*FollowMap).data {
-			if _, ok := followMapA.(*FollowMap).data[following]; ok {
-				intersection = append(intersection, following)
+	}
+
+	// Remove the smallest map from the list of maps to intersect
+	followMaps = followMaps[1:]
+
+	intersection := make([]uint64, 0)
+	for follower := range smallest.data {
+		found := true
+		for _, followMap := range followMaps {
+			if _, ok := followMap.data[follower]; !ok {
+				found = false
+				break
 			}
 		}
+		if found {
+			intersection = append(intersection, follower)
+		}
 	}
 
 	return intersection, nil

diff --git a/graphd/graph_test.go b/graphd/graph_test.go
@@ -128,3 +128,59 @@ func BenchmarkDoesFollow(b *testing.B) {
 		graph.DoesFollow(uids[i], uids[i])
 	}
 }
+
+func TestIntersectNFollowers(t *testing.T) {
+	graph := graphd.NewGraph()
+
+	dids := make([]string, 100)
+
+	// Generate some random DIDs
+	for i := 0; i < 100; i++ {
+		dids[i] = fmt.Sprintf("did:example:%d", i)
+	}
+
+	uids := make([]uint64, 100)
+
+	// Acquire all DIDs
+	for i := 0; i < 100; i++ {
+		uids[i] = graph.AcquireDID(dids[i])
+	}
+
+	// Create a known set of followers
+	for i := 0; i < 100; i++ {
+		graph.AddFollow(uids[i], uids[0])
+	}
+
+	for i := 5; i < 50; i++ {
+		graph.AddFollow(uids[i], uids[1])
+	}
+
+	for i := 2; i < 25; i++ {
+		graph.AddFollow(uids[i], uids[2])
+	}
+
+	for i := 0; i < 10; i++ {
+		graph.AddFollow(uids[i], uids[3])
+	}
+
+	// Intersect the followers of the first 4 DIDs
+	intersect, err := graph.IntersectFollowers(uids[:4])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(intersect) != 5 {
+		t.Errorf("expected 5 followers, got %d", len(intersect))
+	}
+
+	// Check the followers are the expected ones
+	intersectMap := make(map[uint64]bool)
+	for _, uid := range intersect {
+		intersectMap[uid] = true
+	}
+
+	for i := 5; i < 10; i++ {
+		if !intersectMap[uids[i]] {
+			t.Errorf("expected uid %d to be in the intersection", uids[i])
+		}
+	}
+}
diff --git a/graphd/handlers/handlers.go b/graphd/handlers/handlers.go
@@ -146,20 +146,20 @@ func (h *Handlers) GetAreMoots(c echo.Context) error {
 }
 
 func (h *Handlers) GetIntersectFollowers(c echo.Context) error {
-	didA := c.QueryParam("didA")
-	didB := c.QueryParam("didB")
-
-	uidA, ok := h.graph.GetUID(didA)
-	if !ok {
-		return c.JSON(404, "actor uid not found")
-	}
-
-	uidB, ok := h.graph.GetUID(didB)
-	if !ok {
-		return c.JSON(404, "target uid not found")
+	if !c.QueryParams().Has("did") {
+		return c.JSON(400, "did query param is required")
+	}
+	qDIDs := c.QueryParams()["did"]
+	uids := make([]uint64, 0)
+	for _, qDID := range qDIDs {
+		uid, ok := h.graph.GetUID(qDID)
+		if !ok {
+			return c.JSON(404, fmt.Sprintf("uid not found for did %s", qDID))
+		}
+		uids = append(uids, uid)
 	}
 
-	intersect, err := h.graph.IntersectFollowers(uidA, uidB)
+	intersect, err := h.graph.IntersectFollowers(uids)
 	if err != nil {
 		slog.Error("failed to intersect followers", "err", err)
 		return c.JSON(500, "failed to intersect followers")
@@ -175,20 +175,20 @@ func (h *Handlers) GetIntersectFollowers(c echo.Context) error {
 }
 
 func (h *Handlers) GetIntersectFollowing(c echo.Context) error {
-	didA := c.QueryParam("didA")
-	didB := c.QueryParam("didB")
-
-	uidA, ok := h.graph.GetUID(didA)
-	if !ok {
-		return c.JSON(404, "actor uid not found")
-	}
-
-	uidB, ok := h.graph.GetUID(didB)
-	if !ok {
-		return c.JSON(404, "target uid not found")
+	if !c.QueryParams().Has("did") {
+		return c.JSON(400, "did query param is required")
+	}
+	qDIDs := c.QueryParams()["did"]
+	uids := make([]uint64, 0)
+	for _, qDID := range qDIDs {
+		uid, ok := h.graph.GetUID(qDID)
+		if !ok {
+			return c.JSON(404, fmt.Sprintf("uid not found for did %s", qDID))
+		}
+		uids = append(uids, uid)
 	}
 
-	intersect, err := h.graph.IntersectFollowing(uidA, uidB)
+	intersect, err := h.graph.IntersectFollowing(uids)
 	if err != nil {
 		slog.Error("failed to intersect following", "err", err)
 		return c.JSON(500, "failed to intersect following")