introduce option for recursively checking all child rooms as well

roominfotree/roominfotree.go

@@ -97,6 +97,60 @@ func resolveRooms(rooms []string, client *mautrix.Client) ([]id.RoomID, map[id.R
 	return roomIDs, aliasesByRoomID
 }
 
+func getChildRoomChunks(roomID id.RoomID, client *mautrix.Client) []*mautrix.ChildRoomsChunk {
+	var maxDepth *int
+	maxDepth = nil
+	suggestedOnly := false
+
+	var childRoomChunks []*mautrix.ChildRoomsChunk
+
+	hierarchyRequest := mautrix.ReqHierarchy{
+		MaxDepth:      maxDepth,
+		SuggestedOnly: suggestedOnly,
+	}
+	finishedRequesting := false
+	for !finishedRequesting {
+		hierarchyResponse, err := client.Hierarchy(context.Background(), roomID, &hierarchyRequest)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		childRoomChunks = append(childRoomChunks, hierarchyResponse.Rooms...)
+
+		if hierarchyResponse.NextBatch == "" {
+			finishedRequesting = true
+		} else {
+			hierarchyRequest = mautrix.ReqHierarchy{
+				From:          hierarchyResponse.NextBatch,
+				Limit:         hierarchyRequest.Limit,
+				MaxDepth:      hierarchyRequest.MaxDepth,
+				SuggestedOnly: hierarchyRequest.SuggestedOnly,
+			}
+		}
+	}
+
+	return childRoomChunks
+}
+
+func addChildRooms(givenRoomIDs []id.RoomID, client *mautrix.Client) []id.RoomID {
+	roomIDSet := make(map[id.RoomID]bool)
+
+	for _, givenRoomID := range givenRoomIDs {
+		roomIDSet[givenRoomID] = true
+
+		childRoomChunks := getChildRoomChunks(givenRoomID, client)
+		for _, childRoomChunk := range childRoomChunks {
+			roomIDSet[childRoomChunk.PublicRoomInfo.RoomID] = true
+		}
+	}
+
+	roomIDs := make([]id.RoomID, 0, len(roomIDSet))
+	for roomID := range roomIDSet {
+		roomIDs = append(roomIDs, roomID)
+	}
+	return roomIDs
+}
+
 func getMembersByHomeserverByRoomID(roomIDs []id.RoomID, client *mautrix.Client) map[id.RoomID](map[string]uint) {
 	joinedMembersByRoomID := make(map[id.RoomID]*mautrix.RespJoinedMembers)
 	for _, roomID := range roomIDs {
@@ -174,8 +228,11 @@ func getServerVersionInfoByHomeserver(homeservers []string, federationClient *fc
 	return serverVersionInfoByHomeserver
 }
 
-func Get(rooms []string, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
+func Get(rooms []string, recursive bool, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
 	roomIDs, aliasesByRoomID := resolveRooms(rooms, client)
+	if recursive {
+		roomIDs = addChildRooms(roomIDs, client)
+	}
 	membersByHomeserverByRoomID := getMembersByHomeserverByRoomID(roomIDs, client)
 	homeservers := getHomeservers(membersByHomeserverByRoomID)
 	serverVersionInfoByHomeserver := getServerVersionInfoByHomeserver(homeservers, federationClient)