From f805786ae09b2b1c896d208ef4640195433ce7cb Mon Sep 17 00:00:00 2001
From: June <june@jsts.xyz>
Date: Mon, 18 Aug 2025 23:16:54 +0200
Subject: [PATCH] move room ID resolution and given alias information to
 roomInfoTree

Move room ID resolution into roomInfoTree and therefore store given
alias information in RoomInfoTree as well.
---
 main.go                      | 36 +++---------------------
 roominfotree/roominfotree.go | 53 ++++++++++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 34 deletions(-)

diff --git a/main.go b/main.go
index e775f2c..ede617e 100644
--- a/main.go
+++ b/main.go
@@ -73,43 +73,15 @@ func main() {
 		fclient.WithTimeout(config.HomeserverVersionInfoTimeout),
 	)
 
-	roomIDSet := make(map[id.RoomID]bool)
-	aliasSetByRoomID := make(map[id.RoomID](map[string]bool))
-	for _, room := range config.Rooms {
-		// Check, if given room is an alias and try to resolve it into a room id.
-		if strings.HasPrefix(room, "#") {
-			resolvedAlias, err := client.ResolveAlias(context.Background(), id.RoomAlias(room))
-			if err != nil {
-				log.Fatal(err)
-			}
-
-			roomIDSet[resolvedAlias.RoomID] = true
-
-			aliasSet, ok := aliasSetByRoomID[resolvedAlias.RoomID]
-			if !ok {
-				aliasSet = make(map[string]bool)
-				aliasSetByRoomID[resolvedAlias.RoomID] = aliasSet
-			}
-			aliasSet[room] = true
-		} else {
-			roomIDSet[id.RoomID(room)] = true
-		}
-	}
-
-	roomIDs := make([]id.RoomID, 0, len(roomIDSet))
-	for roomID := range roomIDSet {
-		roomIDs = append(roomIDs, roomID)
-	}
-
-	roomInfoTree := roominfotree.Get(roomIDs, client, federationClient)
+	roomInfoTree := roominfotree.Get(config.Rooms, client, federationClient)
 
 	for roomID, roomInfo := range roomInfoTree {
 		fmt.Println("Room:")
 		fmt.Printf("  %s -> %d\n", roomID, roomInfo.MemberCount)
-		aliasSet, ok := aliasSetByRoomID[roomID]
-		if ok {
+		aliases := roomInfo.GivenAliases
+		if len(aliases) > 0 {
 			fmt.Println("Given Aliases:")
-			for alias := range aliasSet {
+			for _, alias := range aliases {
 				fmt.Printf("  %s\n", alias)
 			}
 		}
diff --git a/roominfotree/roominfotree.go b/roominfotree/roominfotree.go
index 185b2f9..96a3b54 100644
--- a/roominfotree/roominfotree.go
+++ b/roominfotree/roominfotree.go
@@ -3,6 +3,7 @@ package roominfotree
 import (
 	"context"
 	"log"
+	"strings"
 	"sync"
 
 	"github.com/matrix-org/gomatrixserverlib/fclient"
@@ -20,6 +21,7 @@ type RoomInfoTree map[id.RoomID](*RoomInfo)
 
 type RoomInfo struct {
 	MemberCount     uint
+	GivenAliases    []string
 	MaxRoomVersions map[string](*MaxRoomVersionInfo)
 }
 
@@ -52,6 +54,49 @@ var unknownServerVersionInfo = fclient.Version{
 	},
 }
 
+// Resolves the given list of rooms to a list of room IDs and a map of given aliases by room ID.
+func resolveRooms(rooms []string, client *mautrix.Client) ([]id.RoomID, map[id.RoomID]([]string)) {
+	roomIDSet := make(map[id.RoomID]bool)
+	aliasSetByRoomID := make(map[id.RoomID](map[string]bool))
+	for _, room := range rooms {
+		// Check, if given room is an alias and try to resolve it into a room id.
+		if strings.HasPrefix(room, "#") {
+			resolvedAlias, err := client.ResolveAlias(context.Background(), id.RoomAlias(room))
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			roomIDSet[resolvedAlias.RoomID] = true
+
+			aliasSet, ok := aliasSetByRoomID[resolvedAlias.RoomID]
+			if !ok {
+				aliasSet = make(map[string]bool)
+				aliasSetByRoomID[resolvedAlias.RoomID] = aliasSet
+			}
+			aliasSet[room] = true
+		} else {
+			roomIDSet[id.RoomID(room)] = true
+		}
+	}
+
+	roomIDs := make([]id.RoomID, 0, len(roomIDSet))
+	for roomID := range roomIDSet {
+		roomIDs = append(roomIDs, roomID)
+	}
+
+	aliasesByRoomID := make(map[id.RoomID]([]string))
+	for roomID, aliasSet := range aliasSetByRoomID {
+		aliases := make([]string, 0, len(aliasSet))
+		for alias := range aliasSet {
+			aliases = append(aliases, alias)
+		}
+
+		aliasesByRoomID[roomID] = aliases
+	}
+
+	return roomIDs, aliasesByRoomID
+}
+
 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 {
@@ -129,7 +174,8 @@ func getServerVersionInfoByHomeserver(homeservers []string, federationClient *fc
 	return serverVersionInfoByHomeserver
 }
 
-func Get(roomIDs []id.RoomID, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
+func Get(rooms []string, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
+	roomIDs, aliasesByRoomID := resolveRooms(rooms, client)
 	membersByHomeserverByRoomID := getMembersByHomeserverByRoomID(roomIDs, client)
 	homeservers := getHomeservers(membersByHomeserverByRoomID)
 	serverVersionInfoByHomeserver := getServerVersionInfoByHomeserver(homeservers, federationClient)
@@ -139,7 +185,10 @@ func Get(roomIDs []id.RoomID, client *mautrix.Client, federationClient *fclient.
 	for roomID, membersByHomeserver := range membersByHomeserverByRoomID {
 		roomInfo, ok := roomInfoTree[roomID]
 		if !ok {
-			roomInfo = &RoomInfo{MaxRoomVersions: make(map[string]*MaxRoomVersionInfo)}
+			roomInfo = &RoomInfo{
+				GivenAliases:    aliasesByRoomID[roomID],
+				MaxRoomVersions: make(map[string]*MaxRoomVersionInfo),
+			}
 			roomInfoTree[roomID] = roomInfo
 		}