<commit_before>// Copyright 2016 The NorthShore Authors All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/Mirantis/northshore/fsm"
	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
)

var demoBlueprintPath string
var demoBp Blueprint
var demoBpPl *fsm.BlueprintPipeline

// demoBlueprintCmd represents the "demo-blueprint" command
var demoBlueprintCmd = &cobra.Command{
	Use:   "demo-blueprint",
	Short: "Run execution of blueprint",
	Long:  `This command read, parse and process blueprint.`,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("Run Blueprint")
		fmt.Printf("PATH -> %s \n", demoBlueprintPath)
		bp, err := ParseBlueprint(demoBlueprintPath)
		if err != nil {
			fmt.Printf("Parsing error: %s \n", err)
			return
		}
		fmt.Printf("BLUEPRINT -> %+v \n", bp)
	},
}

// demoFSMCmd represents the "demo-fsm" command
var demoFSMCmd = &cobra.Command{
	Use:   "demo-fsm",
	Short: "Demo FSM",
	Long:  `Run the Blueprint Pipeline thru states`,
	Run: func(cmd *cobra.Command, args []string) {

		stages := []string{"Stage A", "Stage B"}
		pl := fsm.NewBlueprintPipeline(stages)

		pl.Start()
		pl.Update(map[string]fsm.StageState{"Stage B": fsm.StageStateRunning})
		pl.Update(map[string]fsm.StageState{"Stage A": fsm.StageStateRunning, "Stage B": fsm.StageStatePaused})
		pl.Update(map[string]fsm.StageState{"Stage B": fsm.StageStateRunning})

	},
}

// demoCmd represents the "demo" command
var demoCmd = &cobra.Command{
	Use:   "demo",
	Short: "Run demo",
	Long: `Run demo on local server.

The local server binds localhost:8998.
Demo Blueprint Pipeline goes thru states.`,
	Run: func(cmd *cobra.Command, args []string) {

		/* Run Blueprint */
		log.Println("#run_blueprint")
		log.Printf("PATH -> %s \n", demoBlueprintPath)
		var err error
		demoBp, err = ParseBlueprint(demoBlueprintPath)
		if err != nil {
			log.Printf("Parsing error: %s \n", err)
			return
		}
		log.Printf("BLUEPRINT -> %+v \n", demoBp)

		/* Run States */
		log.Println("#run_states")
		var stages []string
		for k := range demoBp.Stages {
			stages = append(stages, k)
		}

		ss := []map[string]fsm.StageState{
			{stages[0]: fsm.StageStatePaused},
			{stages[0]: fsm.StageStateRunning, stages[1]: fsm.StageStatePaused},
			{stages[1]: fsm.StageStateRunning},
		}

		go func() {
			for {
				demoBpPl = fsm.NewBlueprintPipeline(stages)
				demoBpPl.Start()

				time.Sleep(time.Second * 9)

				for _, s := range stages {
					time.Sleep(time.Second * 3)
					v := map[string]fsm.StageState{s: fsm.StageStateCreated}
					log.Println("#pl-update", v)
					demoBpPl.Update(v)
				}
				for _, s := range stages {
					time.Sleep(time.Second * 3)
					v := map[string]fsm.StageState{s: fsm.StageStateRunning}
					log.Println("#pl-update", v)
					demoBpPl.Update(v)
				}
				for _, v := range ss {
					time.Sleep(time.Second * 3)
					log.Println("#pl-update", v)
					demoBpPl.Update(v)
				}
			}
		}()

		/* Run local server */
		log.Println("#run_local_server")
		r := mux.NewRouter()

		uiAPI1 := r.PathPrefix("/ui/api/v1").Subrouter().StrictSlash(true)
		uiAPI1.HandleFunc("/", uiAPI1RootHandler).Methods("GET")

		uiAPI1.HandleFunc("/action", demouiAPI1ActionHandler).Methods("GET", "POST")
		uiAPI1.HandleFunc("/blueprints", demouiAPI1BlueprintsHandler).Methods("GET", "POST")
		uiAPI1.HandleFunc("/errors", demouiAPI1ErrorsHandler).Methods("GET", "POST")

		ui := r.PathPrefix("/ui").Subrouter().StrictSlash(true)
		ui.PathPrefix("/{uiDir:(app)|(assets)|(node_modules)}").Handler(http.StripPrefix("/ui", http.FileServer(http.Dir("ui/"))))
		ui.HandleFunc("/{s}", uiIndexHandler)
		ui.HandleFunc("/", uiIndexHandler)

		// with 'nshore run local', you can got to http://localhost:8998/ and see a list of
		// what is in static ... if you put index.html in there, it'll be returned.
		// NB: do not put /static in the path, that'll get you a 404.
		r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("static/"))))

		log.Println("Listening at port 8998")
		http.ListenAndServe(":8998", r)

	},
}

func init() {
	demoBlueprintCmd.Flags().StringVarP(&demoBlueprintPath, "file", "f", ".", "Path to blueprint yaml")
	demoCmd.Flags().StringVarP(&demoBlueprintPath, "file", "f", ".", "Path to blueprint yaml")
	runCmd.AddCommand(demoBlueprintCmd)
	runCmd.AddCommand(demoFSMCmd)
	runCmd.AddCommand(demoCmd)
}

func demouiAPI1ActionHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	o := map[string]interface{}{
		"data": []map[string]interface{}{
			{"details": "Details 1"},
			{"details": "Details 2"},
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1ActionHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}

func demouiAPI1BlueprintsHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	type BP struct {
		Blueprint
		*fsm.BlueprintPipeline
	}

	o := map[string]interface{}{
		"data": []BP{
			{demoBp, demoBpPl},
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1BlueprintsHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}

func demouiAPI1ErrorsHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	o := map[string]interface{}{
		"data": []map[string]interface{}{},
		"errors": []map[string]interface{}{
			{
				"details": "Error details 1",
				"title":   "Error title 1",
			},
			{
				"details": "Details of Error 2",
				"meta": map[string]interface{}{
					"info": "meta info of Error 2",
				},
			},
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1ErrorsHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}
<commit_msg>Redefine demoBp<commit_after>// Copyright 2016 The NorthShore Authors All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/Mirantis/northshore/fsm"
	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
)

// BP represents a combined data of the Blueprint with States
// TODO: refactor Blueprint to integrate State info
// the State should be updated on changing the stages via stages setter
type BP struct {
	*Blueprint
	*fsm.BlueprintPipeline
}

var demoBlueprintPath string
var demoBp BP

// demoBlueprintCmd represents the "demo-blueprint" command
var demoBlueprintCmd = &cobra.Command{
	Use:   "demo-blueprint",
	Short: "Run execution of blueprint",
	Long:  `This command read, parse and process blueprint.`,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("Run Blueprint")
		fmt.Printf("PATH -> %s \n", demoBlueprintPath)
		bp, err := ParseBlueprint(demoBlueprintPath)
		if err != nil {
			fmt.Printf("Parsing error: %s \n", err)
			return
		}
		fmt.Printf("BLUEPRINT -> %+v \n", bp)
	},
}

// demoFSMCmd represents the "demo-fsm" command
var demoFSMCmd = &cobra.Command{
	Use:   "demo-fsm",
	Short: "Demo FSM",
	Long:  `Run the Blueprint Pipeline thru states`,
	Run: func(cmd *cobra.Command, args []string) {

		stages := []string{"Stage A", "Stage B"}
		pl := fsm.NewBlueprintPipeline(stages)

		pl.Start()
		pl.Update(map[string]fsm.StageState{"Stage B": fsm.StageStateRunning})
		pl.Update(map[string]fsm.StageState{"Stage A": fsm.StageStateRunning, "Stage B": fsm.StageStatePaused})
		pl.Update(map[string]fsm.StageState{"Stage B": fsm.StageStateRunning})

	},
}

// demoCmd represents the "demo" command
var demoCmd = &cobra.Command{
	Use:   "demo",
	Short: "Run demo",
	Long: `Run demo on local server.

The local server binds localhost:8998.
Demo Blueprint Pipeline goes thru states.`,
	Run: func(cmd *cobra.Command, args []string) {

		/* Run Blueprint */
		log.Println("#run_blueprint")
		log.Printf("PATH -> %s \n", demoBlueprintPath)
		bp, err := ParseBlueprint(demoBlueprintPath)
		if err != nil {
			log.Printf("Parsing error: %s \n", err)
			return
		}
		log.Printf("BLUEPRINT -> %+v \n", demoBp)

		/* Run States */
		log.Println("#run_states")
		var stages []string
		for k := range bp.Stages {
			stages = append(stages, k)
		}

		ss := []map[string]fsm.StageState{
			{stages[0]: fsm.StageStatePaused},
			{stages[0]: fsm.StageStateRunning, stages[1]: fsm.StageStatePaused},
			{stages[1]: fsm.StageStateRunning},
		}

		go func() {
			for {
				pl := fsm.NewBlueprintPipeline(stages)
				demoBp = BP{&bp, pl}
				demoBp.Start()

				time.Sleep(time.Second * 9)

				for _, s := range stages {
					time.Sleep(time.Second * 3)
					v := map[string]fsm.StageState{s: fsm.StageStateCreated}
					log.Println("#pl-update", v)
					demoBp.Update(v)
				}
				for _, s := range stages {
					time.Sleep(time.Second * 3)
					v := map[string]fsm.StageState{s: fsm.StageStateRunning}
					log.Println("#pl-update", v)
					demoBp.Update(v)
				}
				for _, v := range ss {
					time.Sleep(time.Second * 3)
					log.Println("#pl-update", v)
					demoBp.Update(v)
				}
			}
		}()

		/* Run local server */
		log.Println("#run_local_server")
		r := mux.NewRouter()

		uiAPI1 := r.PathPrefix("/ui/api/v1").Subrouter().StrictSlash(true)
		uiAPI1.HandleFunc("/", uiAPI1RootHandler).Methods("GET")

		uiAPI1.HandleFunc("/action", demouiAPI1ActionHandler).Methods("GET", "POST")
		uiAPI1.HandleFunc("/blueprints", demouiAPI1BlueprintsHandler).Methods("GET", "POST")
		uiAPI1.HandleFunc("/errors", demouiAPI1ErrorsHandler).Methods("GET", "POST")

		ui := r.PathPrefix("/ui").Subrouter().StrictSlash(true)
		ui.PathPrefix("/{uiDir:(app)|(assets)|(node_modules)}").Handler(http.StripPrefix("/ui", http.FileServer(http.Dir("ui/"))))
		ui.HandleFunc("/{s}", uiIndexHandler)
		ui.HandleFunc("/", uiIndexHandler)

		// with 'nshore run local', you can got to http://localhost:8998/ and see a list of
		// what is in static ... if you put index.html in there, it'll be returned.
		// NB: do not put /static in the path, that'll get you a 404.
		r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("static/"))))

		log.Println("Listening at port 8998")
		http.ListenAndServe(":8998", r)

	},
}

func init() {
	demoBlueprintCmd.Flags().StringVarP(&demoBlueprintPath, "file", "f", ".", "Path to blueprint yaml")
	demoCmd.Flags().StringVarP(&demoBlueprintPath, "file", "f", ".", "Path to blueprint yaml")
	runCmd.AddCommand(demoBlueprintCmd)
	runCmd.AddCommand(demoFSMCmd)
	runCmd.AddCommand(demoCmd)
}

func demouiAPI1ActionHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	o := map[string]interface{}{
		"data": []map[string]interface{}{
			{"details": "Details 1"},
			{"details": "Details 2"},
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1ActionHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}

func demouiAPI1BlueprintsHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	o := map[string]interface{}{
		"data": []BP{
			demoBp,
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1BlueprintsHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}

func demouiAPI1ErrorsHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/vnd.api+json")

	o := map[string]interface{}{
		"data": []map[string]interface{}{},
		"errors": []map[string]interface{}{
			{
				"details": "Error details 1",
				"title":   "Error title 1",
			},
			{
				"details": "Details of Error 2",
				"meta": map[string]interface{}{
					"info": "meta info of Error 2",
				},
			},
		},
		"meta": map[string]interface{}{
			"info": "demouiAPI1ErrorsHandler",
		},
	}

	json.NewEncoder(w).Encode(o)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/dendrite/setup/base"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
var filterType = flag.String("filtertype", "", "the event types to filter on")

func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	cfg.Logging = append(cfg.Logging[:0], config.LogrusHook{
		Type:  "std",
		Level: "error",
	})
	base := base.NewBaseDendrite(cfg, "ResolveState", base.DisableMetrics)
	args := flag.Args()

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	cache, err := caching.NewInMemoryLRUCache(true)
	if err != nil {
		panic(err)
	}

	roomserverDB, err := storage.Open(base, &cfg.RoomServer.Database, cache)
	if err != nil {
		panic(err)
	}

	blockNIDs, err := roomserverDB.StateBlockNIDs(ctx, snapshotNIDs)
	if err != nil {
		panic(err)
	}

	var stateEntries []types.StateEntryList
	for _, list := range blockNIDs {
		entries, err2 := roomserverDB.StateEntries(ctx, list.StateBlockNIDs)
		if err2 != nil {
			panic(err2)
		}
		stateEntries = append(stateEntries, entries...)
	}

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		for _, e := range entry.StateEntries {
			eventNIDs = append(eventNIDs, e.EventNID)
		}
	}

	fmt.Println("Fetching", len(eventNIDs), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	var resolved Events
	resolved, err = gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	sort.Sort(resolved)
	filteringEventType := *filterType
	count := 0
	for _, event := range resolved {
		if filteringEventType != "" && event.Type() != filteringEventType {
			continue
		}
		count++
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}

	fmt.Println()
	fmt.Println("Returned", count, "state events after filtering")
}

type Events []*gomatrixserverlib.Event

func (e Events) Len() int {
	return len(e)
}

func (e Events) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e Events) Less(i, j int) bool {
	typeDelta := strings.Compare(e[i].Type(), e[j].Type())
	if typeDelta < 0 {
		return true
	}
	if typeDelta > 0 {
		return false
	}
	stateKeyDelta := strings.Compare(*e[i].StateKey(), *e[j].StateKey())
	return stateKeyDelta < 0
}
<commit_msg>Fix bugs in `resolve-state` tool<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/matrix-org/dendrite/internal/caching"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup"
	"github.com/matrix-org/dendrite/setup/base"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
)

// This is a utility for inspecting state snapshots and running state resolution
// against real snapshots in an actual database.
// It takes one or more state snapshot NIDs as arguments, along with a room version
// to use for unmarshalling events, and will produce resolved output.
//
// Usage: ./resolve-state --roomversion=version snapshot [snapshot ...]
//   e.g. ./resolve-state --roomversion=5 1254 1235 1282

var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
var filterType = flag.String("filtertype", "", "the event types to filter on")

func main() {
	ctx := context.Background()
	cfg := setup.ParseFlags(true)
	cfg.Logging = append(cfg.Logging[:0], config.LogrusHook{
		Type:  "std",
		Level: "error",
	})
	base := base.NewBaseDendrite(cfg, "ResolveState", base.DisableMetrics)
	args := flag.Args()

	fmt.Println("Room version", *roomVersion)

	snapshotNIDs := []types.StateSnapshotNID{}
	for _, arg := range args {
		if i, err := strconv.Atoi(arg); err == nil {
			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
		}
	}

	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")

	cache, err := caching.NewInMemoryLRUCache(true)
	if err != nil {
		panic(err)
	}

	roomserverDB, err := storage.Open(base, &cfg.RoomServer.Database, cache)
	if err != nil {
		panic(err)
	}

	stateres := state.NewStateResolution(roomserverDB, &types.RoomInfo{
		RoomVersion: gomatrixserverlib.RoomVersion(*roomVersion),
	})

	var stateEntries []types.StateEntry
	for _, snapshotNID := range snapshotNIDs {
		var entries []types.StateEntry
		entries, err = stateres.LoadStateAtSnapshot(ctx, snapshotNID)
		if err != nil {
			panic(err)
		}
		stateEntries = append(stateEntries, entries...)
	}

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		eventNIDs = append(eventNIDs, entry.EventNID)
	}

	fmt.Println("Fetching", len(eventNIDs), "state events")
	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
	if err != nil {
		panic(err)
	}

	authEventIDMap := make(map[string]struct{})
	events := make([]*gomatrixserverlib.Event, len(eventEntries))
	for i := range eventEntries {
		events[i] = eventEntries[i].Event
		for _, authEventID := range eventEntries[i].AuthEventIDs() {
			authEventIDMap[authEventID] = struct{}{}
		}
	}

	authEventIDs := make([]string, 0, len(authEventIDMap))
	for authEventID := range authEventIDMap {
		authEventIDs = append(authEventIDs, authEventID)
	}

	fmt.Println("Fetching", len(authEventIDs), "auth events")
	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
	if err != nil {
		panic(err)
	}

	authEvents := make([]*gomatrixserverlib.Event, len(authEventEntries))
	for i := range authEventEntries {
		authEvents[i] = authEventEntries[i].Event
	}

	fmt.Println("Resolving state")
	var resolved Events
	resolved, err = gomatrixserverlib.ResolveConflicts(
		gomatrixserverlib.RoomVersion(*roomVersion),
		events,
		authEvents,
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("Resolved state contains", len(resolved), "events")
	sort.Sort(resolved)
	filteringEventType := *filterType
	count := 0
	for _, event := range resolved {
		if filteringEventType != "" && event.Type() != filteringEventType {
			continue
		}
		count++
		fmt.Println()
		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
		fmt.Printf("  %s\n", string(event.Content()))
	}

	fmt.Println()
	fmt.Println("Returned", count, "state events after filtering")
}

type Events []*gomatrixserverlib.Event

func (e Events) Len() int {
	return len(e)
}

func (e Events) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e Events) Less(i, j int) bool {
	typeDelta := strings.Compare(e[i].Type(), e[j].Type())
	if typeDelta < 0 {
		return true
	}
	if typeDelta > 0 {
		return false
	}
	stateKeyDelta := strings.Compare(*e[i].StateKey(), *e[j].StateKey())
	return stateKeyDelta < 0
}
<|endoftext|>


<commit_before>/*
  A partitioner assigns partitions within a consumer group like a
  consistent hash would. When group membership changes the least
  number of partitions are reassigned.

  I don't actually use a consistent hash because such algos take
  either a lot of CPU (the max(N-hash-functions) algo) or a lot
  of ram (the points-on-the-circle algo), or aren't applicable
  when members come and go (jump-hash).

  Instead I take advantage of the fact that I can know the current
  assignments of the remaining group members when constructing
  the new assignments. That makes keep the partitioning stable
  relatively easy to do:
    1) calculate an ideal load for each consumer (I assume equal
	   weight since I don't need anything more complicated...yet)
	2) remove excess partitions from overloaded consumers
	3) add the abandonned and excess partitions to the least loaded
	   consumers.

  Note that the state of the partition at time T depends on
  the history of the consumer group memberships from the time when
  there was 1 consumer until time T.

  Copyright 2016 MistSys
*/

package stable

import (
	"fmt"
	"log"
	"sort"

	"github.com/Shopify/sarama"
)

// a partitioner that assigns partitions to consumers such that as consumers come and go the least number of partitions are reassigned
type stablePartitioner string

// global instance of the stable partitioner
const Stable stablePartitioner = "stable"

// print a debug message
func dbgf(format string, args ...interface{}) {
	log.Printf(format, args...)
}

func (sp stablePartitioner) PrepareJoin(jreq *sarama.JoinGroupRequest, topics []string, current_assignments map[string][]int32) {
	// encode the current assignments in a manner proprietary to this partitioner
	var data = data{
		version:     1,
		assignments: current_assignments,
	}

	jreq.AddGroupProtocolMetadata(string(sp),
		&sarama.ConsumerGroupMemberMetadata{
			Version:  1,
			Topics:   topics,
			UserData: data.marshal(),
		})
}

// for each topic in jresp, assign the topic's partitions to the members requesting the topic
func (sp stablePartitioner) Partition(sreq *sarama.SyncGroupRequest, jresp *sarama.JoinGroupResponse, client sarama.Client) error {
	if jresp.GroupProtocol != string(sp) {
		return fmt.Errorf("sarama.JoinGroupResponse.GroupProtocol %q unexpected; expected %q", jresp.GroupProtocol != string(sp))
	}
	by_member, err := jresp.GetMembers() //  map of member to ConsumerGroupMemberMetadata
	dbgf("by_member = %v", by_member)
	if err != nil {
		return err
	}

	// invert the data, so we have the requests grouped by topic (they arrived grouped by member, since the kafka broker treats the data from each consumer as an opaque blob, so it couldn't do this step for us)
	by_topic := make(map[string]map[string][]int32) // map of topic to members and members to current partition assignment
	for member, request := range by_member {
		if request.Version != 1 {
			// skip unsupported versions. we'll only assign to clients we can understand. Since we are such a client
			// we won't block all consumers (at least for those topics we consume). If this ends up a bad idea, we
			// can always change this code to return an error.
			continue
		}
		var data data
		err := data.unmarshal(request.UserData)
		if err != nil {
			return err
		}

		for _, topic := range request.Topics {
			members, ok := by_topic[topic]
			if !ok {
				members = make(map[string][]int32)
				by_topic[topic] = members
			}
			members[member] = data.assignments[topic] // NOTE: might be nil, which is OK. It just means the member wants to consume the partition but isn't doing so currently
		}
	}
	dbgf("by_topic = %v", by_topic)

	// lookup the partitions in each topic. since we are asking for all partitions, not just the online ones, the numbering
	// appears to always be 0...N-1. But in case I don't understand kafka and there is some corner case where the numbering
	// of partitions is different I keep careful track of the exact numbers I've received.
	var partitions_by_topic = make(map[string]partitionslist)
	for topic := range by_topic {
		// note: calls to client.Partitions() hit the metadata cache in the sarama client, so we don't gain much by asking concurrently
		partitions, err := client.Partitions(topic)
		if err != nil {
			// what to do? we could maybe skip the topic, assigning it to no-one. But I/O errors are likely to happen again.
			// so let's stop partitioning and return the error.
			return err
		}
		pl := make(partitionslist, len(partitions)) // we must make a copy before we sort the list
		for i, p := range partitions {
			pl[i] = p
		}
		sort.Sort(pl)
		partitions_by_topic[topic] = pl
	}
	dbgf("partitions_by_topic = %v", partitions_by_topic)

	// and compute the sorted set of members of each topic
	var members_by_topic = make(map[string]memberslist)
	for topic, members := range by_topic {
		ml := make(memberslist, 0, len(members))
		for m := range members {
			ml = append(ml, m)
		}
		sort.Sort(ml)
		members_by_topic[topic] = ml
	}
	dbgf("members_by_topic = %v", members_by_topic)

	// I want topics with the same # of partitions and the same consumer group membership to result in the same partition assignments.
	// That way messages published under identical partition keys in those topics will all end up consumed by the same member.
	// So organize topics into groups which will be partitioned identically
	var matched_topics = make(map[string]string) // map from each topic to the 'master' topic with the same # of partitions and group membership. Topics which are unique are their own master
topic_match_loop:
	for topic, members := range members_by_topic {
		// see if a match exists
		num_partitions := len(partitions_by_topic[topic])
		for t := range matched_topics { // TODO if # of topics gets large enough this shows up in the profiler, change this to some sort of a map lookup, rather than this O(N^2) search
			if num_partitions == len(partitions_by_topic[t]) && members.Equal(members_by_topic[t]) {
				// match; have topic 'topic' be partitioned the same way as topic 't'
				matched_topics[topic] = matched_topics[t]
				continue topic_match_loop
			}
		}
		// no existing topic matches this one, so it is its own match
		matched_topics[topic] = topic
	}
	dbgf("matched_topics = %v", matched_topics)

	// adjust the partitioning of each master topic in by_topic
	for topic, match := range matched_topics {
		if topic == match {
			adjust_partitioning(by_topic[topic], partitions_by_topic[topic])
		} // else it is not a master topic. once the master has been partitioned we'll simply copy the result
	}

	// set matched topics to the same assignment as their master
	for topic, match := range matched_topics {
		if topic != match {
			by_topic[topic] = by_topic[match]
		}
	}

	// invert by_topic into the equivalent organized by member, and then by topic
	assignments := make(map[string]map[string][]int32) // map of member to topics, and topic to partitions
	for topic, members := range by_topic {
		for member, partitions := range members {
			topics, ok := assignments[member]
			if !ok {
				topics = make(map[string][]int32)
				assignments[member] = topics
			}
			topics[topic] = partitions
		}
	}
	dbgf("assignments = %v", assignments)

	// and encode the assignments in the sync request
	for member, topics := range assignments {
		sreq.AddGroupAssignmentMember(member,
			&sarama.ConsumerGroupMemberAssignment{
				Version: 1,
				Topics:  topics,
			})
	}

	return nil
}

func (stablePartitioner) ParseSync(sresp *sarama.SyncGroupResponse) (map[string][]int32, error) {
	if len(sresp.MemberAssignment) == 0 {
		// in the corner case that we ask for no topics, we get nothing back. However sarama fd498173ae2bf (head of master branch Nov 6th 2016) will return a useless error if we call sresp.GetMemberAssignment() in this case
		return nil, nil
	}
	ma, err := sresp.GetMemberAssignment()
	dbgf("MemberAssignment %v", ma)
	if err != nil {
		return nil, err
	}
	if ma.Version != 1 {
		return nil, fmt.Errorf("unsupported MemberAssignment version %d", ma.Version)
	}
	return ma.Topics, nil
}

// ----------------------------------

// adjust_partitioning does the main work. it adjusts the partition assignment map it is passed in-place
func adjust_partitioning(assignment map[string][]int32, partitions partitionslist) {
	dbgf("adjust_partitioning(assignment = %v, partitions = %v)", assignment, partitions)
	num_members := len(assignment)
	dbgf("num_members = %v", num_members)
	if num_members == 0 {
		// no one wants this topic; stop now before we /0
		return
	}
	num_partitions := len(partitions)
	dbgf("num_partitions = %v", num_partitions)
	level := (num_partitions + num_members - 1) / num_members // the maximum # of partitions each member should receive
	dbgf("level = %v", level)

	unassigned := make(map[int32]struct{}) // set of unassigned partition ids

	// initially all partitions are unassigned
	for _, p := range partitions {
		unassigned[p] = struct{}{}
	}
	dbgf("unassigned = %v", unassigned)

	// let each member keep up to 'level' of its current assignment
	for m, a := range assignment {
		if len(a) > level {
			a = a[:level]
			assignment[m] = a
		}
		for _, p := range a {
			delete(unassigned, p)
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// assign the unassigned partitions to any member with < level partitions
	for p := range unassigned {
	assignment_loop:
		for m, a := range assignment {
			if len(a) < level {
				a = append(a, p)
				assignment[m] = a
				break assignment_loop
			}
		}
	}

	// and we're done
	dbgf("assignment = %v", assignment)
}

// ----------------------------------

// a sortable, comparible list of members
type memberslist []string // a list of the member ids

// implement sort.Interface
func (ml memberslist) Len() int           { return len(ml) }
func (ml memberslist) Swap(i, j int)      { ml[i], ml[j] = ml[j], ml[i] }
func (ml memberslist) Less(i, j int) bool { return ml[i] < ml[j] }

// and compare two sorted memberslist for equality
func (ml memberslist) Equal(ml2 memberslist) bool {
	if len(ml) != len(ml2) {
		return false
	}
	for i, m := range ml {
		if m != ml2[i] {
			return false
		}
	}
	return true
}

// ----------------------------------

// a sortable, comparible list of partition ids
type partitionslist []int32 // a list of the partition ids

// implement sort.Interface
func (pl partitionslist) Len() int           { return len(pl) }
func (pl partitionslist) Swap(i, j int)      { pl[i], pl[j] = pl[j], pl[i] }
func (pl partitionslist) Less(i, j int) bool { return pl[i] < pl[j] }

// and compare two sorted partitionslist for equality
func (pl partitionslist) Equal(pl2 partitionslist) bool {
	if len(pl) != len(pl2) {
		return false
	}
	for i, m := range pl {
		if m != pl2[i] {
			return false
		}
	}
	return true
}
<commit_msg>clean up some edge cases in adjust_partitioning()<commit_after>/*
  A partitioner assigns partitions within a consumer group like a
  consistent hash would. When group membership changes the least
  number of partitions are reassigned.

  I don't actually use a consistent hash because such algos take
  either a lot of CPU (the max(N-hash-functions) algo) or a lot
  of ram (the points-on-the-circle algo), or aren't applicable
  when members come and go (jump-hash).

  Instead I take advantage of the fact that I can know the current
  assignments of the remaining group members when constructing
  the new assignments. That makes keep the partitioning stable
  relatively easy to do:
    1) calculate an ideal load for each consumer (I assume equal
	   weight since I don't need anything more complicated...yet)
	2) remove excess partitions from overloaded consumers
	3) add the abandonned and excess partitions to the least loaded
	   consumers.

  Note that the state of the partition at time T depends on
  the history of the consumer group memberships from the time when
  there was 1 consumer until time T.

  Copyright 2016 MistSys
*/

package stable

import (
	"fmt"
	"log"
	"sort"

	"github.com/Shopify/sarama"
)

// a partitioner that assigns partitions to consumers such that as consumers come and go the least number of partitions are reassigned
type stablePartitioner string

// global instance of the stable partitioner
const Stable stablePartitioner = "stable"

// print a debug message
func dbgf(format string, args ...interface{}) {
	log.Printf(format, args...)
}

func (sp stablePartitioner) PrepareJoin(jreq *sarama.JoinGroupRequest, topics []string, current_assignments map[string][]int32) {
	// encode the current assignments in a manner proprietary to this partitioner
	var data = data{
		version:     1,
		assignments: current_assignments,
	}

	jreq.AddGroupProtocolMetadata(string(sp),
		&sarama.ConsumerGroupMemberMetadata{
			Version:  1,
			Topics:   topics,
			UserData: data.marshal(),
		})
}

// for each topic in jresp, assign the topic's partitions to the members requesting the topic
func (sp stablePartitioner) Partition(sreq *sarama.SyncGroupRequest, jresp *sarama.JoinGroupResponse, client sarama.Client) error {
	if jresp.GroupProtocol != string(sp) {
		return fmt.Errorf("sarama.JoinGroupResponse.GroupProtocol %q unexpected; expected %q", jresp.GroupProtocol != string(sp))
	}
	by_member, err := jresp.GetMembers() //  map of member to ConsumerGroupMemberMetadata
	dbgf("by_member = %v", by_member)
	if err != nil {
		return err
	}

	// invert the data, so we have the requests grouped by topic (they arrived grouped by member, since the kafka broker treats the data from each consumer as an opaque blob, so it couldn't do this step for us)
	by_topic := make(map[string]map[string][]int32) // map of topic to members and members to current partition assignment
	for member, request := range by_member {
		if request.Version != 1 {
			// skip unsupported versions. we'll only assign to clients we can understand. Since we are such a client
			// we won't block all consumers (at least for those topics we consume). If this ends up a bad idea, we
			// can always change this code to return an error.
			continue
		}
		var data data
		err := data.unmarshal(request.UserData)
		if err != nil {
			return err
		}

		for _, topic := range request.Topics {
			members, ok := by_topic[topic]
			if !ok {
				members = make(map[string][]int32)
				by_topic[topic] = members
			}
			members[member] = data.assignments[topic] // NOTE: might be nil, which is OK. It just means the member wants to consume the partition but isn't doing so currently
		}
	}
	dbgf("by_topic = %v", by_topic)

	// lookup the partitions in each topic. since we are asking for all partitions, not just the online ones, the numbering
	// appears to always be 0...N-1. But in case I don't understand kafka and there is some corner case where the numbering
	// of partitions is different I keep careful track of the exact numbers I've received.
	var partitions_by_topic = make(map[string]partitionslist)
	for topic := range by_topic {
		// note: calls to client.Partitions() hit the metadata cache in the sarama client, so we don't gain much by asking concurrently
		partitions, err := client.Partitions(topic)
		if err != nil {
			// what to do? we could maybe skip the topic, assigning it to no-one. But I/O errors are likely to happen again.
			// so let's stop partitioning and return the error.
			return err
		}
		pl := make(partitionslist, len(partitions)) // we must make a copy before we sort the list
		for i, p := range partitions {
			pl[i] = p
		}
		sort.Sort(pl)
		partitions_by_topic[topic] = pl
	}
	dbgf("partitions_by_topic = %v", partitions_by_topic)

	// and compute the sorted set of members of each topic
	var members_by_topic = make(map[string]memberslist)
	for topic, members := range by_topic {
		ml := make(memberslist, 0, len(members))
		for m := range members {
			ml = append(ml, m)
		}
		sort.Sort(ml)
		members_by_topic[topic] = ml
	}
	dbgf("members_by_topic = %v", members_by_topic)

	// I want topics with the same # of partitions and the same consumer group membership to result in the same partition assignments.
	// That way messages published under identical partition keys in those topics will all end up consumed by the same member.
	// So organize topics into groups which will be partitioned identically
	var matched_topics = make(map[string]string) // map from each topic to the 'master' topic with the same # of partitions and group membership. Topics which are unique are their own master
topic_match_loop:
	for topic, members := range members_by_topic {
		// see if a match exists
		num_partitions := len(partitions_by_topic[topic])
		for t := range matched_topics { // TODO if # of topics gets large enough this shows up in the profiler, change this to some sort of a map lookup, rather than this O(N^2) search
			if num_partitions == len(partitions_by_topic[t]) && members.Equal(members_by_topic[t]) {
				// match; have topic 'topic' be partitioned the same way as topic 't'
				matched_topics[topic] = matched_topics[t]
				continue topic_match_loop
			}
		}
		// no existing topic matches this one, so it is its own match
		matched_topics[topic] = topic
	}
	dbgf("matched_topics = %v", matched_topics)

	// adjust the partitioning of each master topic in by_topic
	for topic, match := range matched_topics {
		if topic == match {
			adjust_partitioning(by_topic[topic], partitions_by_topic[topic])
		} // else it is not a master topic. once the master has been partitioned we'll simply copy the result
	}

	// set matched topics to the same assignment as their master
	for topic, match := range matched_topics {
		if topic != match {
			by_topic[topic] = by_topic[match]
		}
	}

	// invert by_topic into the equivalent organized by member, and then by topic
	assignments := make(map[string]map[string][]int32) // map of member to topics, and topic to partitions
	for topic, members := range by_topic {
		for member, partitions := range members {
			topics, ok := assignments[member]
			if !ok {
				topics = make(map[string][]int32)
				assignments[member] = topics
			}
			topics[topic] = partitions
		}
	}
	dbgf("assignments = %v", assignments)

	// and encode the assignments in the sync request
	for member, topics := range assignments {
		sreq.AddGroupAssignmentMember(member,
			&sarama.ConsumerGroupMemberAssignment{
				Version: 1,
				Topics:  topics,
			})
	}

	return nil
}

func (stablePartitioner) ParseSync(sresp *sarama.SyncGroupResponse) (map[string][]int32, error) {
	if len(sresp.MemberAssignment) == 0 {
		// in the corner case that we ask for no topics, we get nothing back. However sarama fd498173ae2bf (head of master branch Nov 6th 2016) will return a useless error if we call sresp.GetMemberAssignment() in this case
		return nil, nil
	}
	ma, err := sresp.GetMemberAssignment()
	dbgf("MemberAssignment %v", ma)
	if err != nil {
		return nil, err
	}
	if ma.Version != 1 {
		return nil, fmt.Errorf("unsupported MemberAssignment version %d", ma.Version)
	}
	return ma.Topics, nil
}

// ----------------------------------

// adjust_partitioning does the main work. it adjusts the partition assignment map it is passed in-place
func adjust_partitioning(assignment map[string][]int32, partitions partitionslist) {
	dbgf("adjust_partitioning(assignment = %v, partitions = %v)", assignment, partitions)
	num_members := len(assignment)
	dbgf("num_members = %v", num_members)
	if num_members == 0 {
		// no one wants this topic; stop now before we /0
		return
	}
	num_partitions := len(partitions)
	dbgf("num_partitions = %v", num_partitions)
	low := num_partitions / num_members                      // the minimum # of partitions each member should receive
	high := (num_partitions + num_members - 1) / num_members // the maximum # of partitions each member should receive
	dbgf("low = %v, high = %v", low, high)

	unassigned := make(map[int32]struct{}) // set of unassigned partition ids

	// initially all partitions are unassigned
	for _, p := range partitions {
		unassigned[p] = struct{}{}
	}
	dbgf("unassigned = %v", unassigned)

	// let each member keep up to 'high' of its current assignment
	for m, a := range assignment {
		if len(a) > high {
			a = a[:high]
			assignment[m] = a
		}
		for _, p := range a {
			delete(unassigned, p)
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// assign the unassigned partitions to any member with < low partitions
	for p := range unassigned {
	unassigned_loop:
		for m, a := range assignment {
			if len(a) < low {
				a = append(a, p)
				assignment[m] = a
				delete(unassigned, p)
				break unassigned_loop
			}
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// take partitions from any member with > low partitions to give
	// to any member with < low partitions
	for m, a := range assignment {
	stealing_from_the_numerous:
		for len(a) < low {
			for m2, a2 := range assignment {
				n := len(a2)
				if n > low {
					// take the last partition from m2 and give it to m
					assignment[m2] = a2[:n-1]
					a = append(a, a2[n-1])
					assignment[m] = a
					continue stealing_from_the_numerous
				}
			}
		}
	}

	// and finally assign any remaining unassigned partitions to any member with < high partitions
	for p := range unassigned {
	assignment_loop:
		for m, a := range assignment {
			if len(a) < high {
				a = append(a, p)
				assignment[m] = a
				delete(unassigned, p) // not really necessary, since we don't reuse unassigned afterwards
				break assignment_loop
			}
		}
	}
	dbgf("assignment = %v", assignment)
	dbgf("unassigned = %v", unassigned)

	// and we're done
	dbgf("assignment = %v", assignment)
}

// ----------------------------------

// a sortable, comparible list of members
type memberslist []string // a list of the member ids

// implement sort.Interface
func (ml memberslist) Len() int           { return len(ml) }
func (ml memberslist) Swap(i, j int)      { ml[i], ml[j] = ml[j], ml[i] }
func (ml memberslist) Less(i, j int) bool { return ml[i] < ml[j] }

// and compare two sorted memberslist for equality
func (ml memberslist) Equal(ml2 memberslist) bool {
	if len(ml) != len(ml2) {
		return false
	}
	for i, m := range ml {
		if m != ml2[i] {
			return false
		}
	}
	return true
}

// ----------------------------------

// a sortable, comparible list of partition ids
type partitionslist []int32 // a list of the partition ids

// implement sort.Interface
func (pl partitionslist) Len() int           { return len(pl) }
func (pl partitionslist) Swap(i, j int)      { pl[i], pl[j] = pl[j], pl[i] }
func (pl partitionslist) Less(i, j int) bool { return pl[i] < pl[j] }

// and compare two sorted partitionslist for equality
func (pl partitionslist) Equal(pl2 partitionslist) bool {
	if len(pl) != len(pl2) {
		return false
	}
	for i, m := range pl {
		if m != pl2[i] {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func TestMain(m *testing.M) {
	os.MkdirAll("tmp", os.ModeDir|os.ModePerm)
	err := exec.Command("go", "build", "-o", "tmp/sqlfmt").Run()
	if err != nil {
		fmt.Println("Failed to build sqlfmt binary:", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

func sqlfmt(sql string, args ...string) (string, error) {
	cmd := exec.Command("tmp/sqlfmt", args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StdinPipe failed: %v", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StdoutPipe failed: %v", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StderrPipe failed: %v", err)
	}

	err = cmd.Start()
	if err != nil {
		return "", fmt.Errorf("cmd.Start failed: %v", err)
	}

	_, err = fmt.Fprint(stdin, sql)
	if err != nil {
		return "", fmt.Errorf("fmt.Fprint failed: %v", err)
	}

	err = stdin.Close()
	if err != nil {
		return "", fmt.Errorf("stdin.Close failed: %v", err)
	}

	output, err := ioutil.ReadAll(stdout)
	if err != nil {
		return "", fmt.Errorf("ioutil.ReadAll(stdout) failed: %v", err)
	}

	errout, err := ioutil.ReadAll(stderr)
	if err != nil {
		return "", fmt.Errorf("ioutil.ReadAll(stderr) failed: %v", err)
	}

	err = cmd.Wait()
	if err != nil {
		return "", fmt.Errorf("cmd.Wait failed: %v\n%s", err, string(errout))
	}

	return string(output), nil
}

func TestSqlFmt(t *testing.T) {
	tests := []struct {
		inputFile          string
		expectedOutputFile string
	}{
		{
			inputFile:          "simple_select_without_from.sql",
			expectedOutputFile: "simple_select_without_from.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_from.sql",
			expectedOutputFile: "simple_select_with_from.fmt.sql",
		},
		{
			inputFile:          "select_from_aliased.sql",
			expectedOutputFile: "select_from_aliased.fmt.sql",
		},
	}

	for i, tt := range tests {
		input, err := ioutil.ReadFile(path.Join("../../testdata", tt.inputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		expected, err := ioutil.ReadFile(path.Join("../../testdata", tt.expectedOutputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		output, err := sqlfmt(string(input))
		if err != nil {
			t.Errorf("%d. sqlfmt failed with %s: %v", i, tt.inputFile, err)
			continue
		}

		if output != string(expected) {
			actualFileName := path.Join("tmp", fmt.Sprintf("%d.sql", i))
			err = ioutil.WriteFile(actualFileName, []byte(output), os.ModePerm)
			if err != nil {
				t.Fatal(err)
			}

			t.Errorf("%d. Given %s, did not receive %s. Unexpected output written to %s", i, tt.inputFile, tt.expectedOutputFile, actualFileName)
		}
	}
}
<commit_msg>Test with []byte instead of strings<commit_after>package main_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func TestMain(m *testing.M) {
	os.MkdirAll("tmp", os.ModeDir|os.ModePerm)
	err := exec.Command("go", "build", "-o", "tmp/sqlfmt").Run()
	if err != nil {
		fmt.Println("Failed to build sqlfmt binary:", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

func sqlfmt(sql []byte, args ...string) ([]byte, error) {
	cmd := exec.Command("tmp/sqlfmt", args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("cmd.StdinPipe failed: %v", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("cmd.StdoutPipe failed: %v", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("cmd.StderrPipe failed: %v", err)
	}

	err = cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("cmd.Start failed: %v", err)
	}

	_, err = stdin.Write(sql)
	if err != nil {
		return nil, fmt.Errorf("stdin.Write failed: %v", err)
	}

	err = stdin.Close()
	if err != nil {
		return nil, fmt.Errorf("stdin.Close failed: %v", err)
	}

	output, err := ioutil.ReadAll(stdout)
	if err != nil {
		return nil, fmt.Errorf("ioutil.ReadAll(stdout) failed: %v", err)
	}

	errout, err := ioutil.ReadAll(stderr)
	if err != nil {
		return nil, fmt.Errorf("ioutil.ReadAll(stderr) failed: %v", err)
	}

	err = cmd.Wait()
	if err != nil {
		return nil, fmt.Errorf("cmd.Wait failed: %v\n%s", err, string(errout))
	}

	return output, nil
}

func TestSqlFmt(t *testing.T) {
	tests := []struct {
		inputFile          string
		expectedOutputFile string
	}{
		{
			inputFile:          "simple_select_without_from.sql",
			expectedOutputFile: "simple_select_without_from.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_from.sql",
			expectedOutputFile: "simple_select_with_from.fmt.sql",
		},
		{
			inputFile:          "select_from_aliased.sql",
			expectedOutputFile: "select_from_aliased.fmt.sql",
		},
	}

	for i, tt := range tests {
		input, err := ioutil.ReadFile(path.Join("../../testdata", tt.inputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		expected, err := ioutil.ReadFile(path.Join("../../testdata", tt.expectedOutputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		output, err := sqlfmt(input)
		if err != nil {
			t.Errorf("%d. sqlfmt failed with %s: %v", i, tt.inputFile, err)
			continue
		}

		if bytes.Compare(output, expected) != 0 {
			actualFileName := path.Join("tmp", fmt.Sprintf("%d.sql", i))
			err = ioutil.WriteFile(actualFileName, output, os.ModePerm)
			if err != nil {
				t.Fatal(err)
			}

			t.Errorf("%d. Given %s, did not receive %s. Unexpected output written to %s", i, tt.inputFile, tt.expectedOutputFile, actualFileName)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"

	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

func parseMoves(spec [][2]string) [][2]*tak.Move {
	var out [][2]*tak.Move
	for _, r := range spec {
		var o [2]*tak.Move
		for i, n := range r {
			if n == "" {
				continue
			}
			m, e := ptn.ParseMove(n)
			if e != nil {
				panic("bad ptn")
			}
			o[i] = &m
		}
		out = append(out, o)
	}
	return out
}

func appendMove(transcript []Expectation,
	id string, tm int,
	move [2]*tak.Move) []Expectation {
	transcript = append(transcript, Expectation{
		recv: []string{
			fmt.Sprintf("Game#%s %s", id, playtak.FormatServer(move[0])),
		},
	})
	if move[1] == nil {
		return transcript
	}
	transcript = append(transcript, Expectation{
		send: []string{
			fmt.Sprintf("Game#%s %s", id, playtak.FormatServer(move[1])),
			fmt.Sprintf("Game#%s Time %d %d", id, tm, tm),
		},
	})
	return transcript
}

func TestBasicGame(t *testing.T) {
	moves := parseMoves([][2]string{
		{"a1", "e1"},
		{"e3", "b1"},
		{"e2", "b2"},
		{"Ce4", "a2"},
		{"e5", ""},
	})
	bot := &TestBotStatic{}
	for _, r := range moves {
		bot.moves = append(bot.moves, *r[0])
	}

	startLine := "Game Start 100 5 Taktician vs HonestJoe white 600"
	var transcript []Expectation
	tm := 600
	for _, r := range moves {
		transcript = appendMove(
			transcript, "100", tm, r)
		tm -= 10
	}
	transcript = append(transcript, Expectation{
		send: []string{
			"Game#100 Over R-0",
		},
	})

	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}

func TestUndoGame(t *testing.T) {
	moves := parseMoves([][2]string{
		{"a1", "e1"},
		{"e3", "b1"},
		{"e2", "b2"},
		{"Ce4", "a2"},
		{"e5", ""},
	})
	bot := &TestBotUndo{}
	for _, r := range moves {
		bot.moves = append(bot.moves, *r[0])
	}
	bot.undoPly = 5

	startLine := "Game Start 100 5 Taktician vs HonestJoe white 600"
	var transcript []Expectation
	tm := 600
	for i, r := range moves {
		transcript = appendMove(
			transcript, "100", tm, r)
		tm -= 10
		if i == 2 {
			e := transcript[len(transcript)-1]
			transcript = append(transcript,
				Expectation{
					send: []string{
						"Game#100 RequestUndo",
					},
					recv: []string{
						"Game#100 RequestUndo",
					},
				},
				Expectation{
					send: []string{
						"Game#100 Undo",
					},
				},
				Expectation{
					send: e.send,
				},
			)
		}
	}
	transcript = append(transcript, Expectation{
		send: []string{
			"Game#100 Over R-0",
		},
	})

	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}

func TestThinker(t *testing.T) {
	moves := parseMoves([][2]string{
		{"a1", "e1"},
		{"e3", "b1"},
		{"e2", "b2"},
		{"Ce4", "a2"},
		{"e5", ""},
	})
	bot := &TestBotThinker{}
	for _, r := range moves {
		bot.moves = append(bot.moves, *r[0])
	}

	startLine := "Game Start 100 5 Taktician vs HonestJoe white 600"
	var transcript []Expectation
	tm := 600
	for _, r := range moves {
		transcript = appendMove(
			transcript, "100", tm, r)
		tm -= 10
	}
	transcript = append(transcript, Expectation{
		send: []string{
			"Game#100 Over R-0",
		},
	})

	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}
<commit_msg>DRY tests a bit<commit_after>package main

import (
	"fmt"
	"testing"

	"github.com/nelhage/taktician/playtak"
	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

func parseMoves(spec [][2]string) [][2]*tak.Move {
	var out [][2]*tak.Move
	for _, r := range spec {
		var o [2]*tak.Move
		for i, n := range r {
			if n == "" {
				continue
			}
			m, e := ptn.ParseMove(n)
			if e != nil {
				panic("bad ptn")
			}
			o[i] = &m
		}
		out = append(out, o)
	}
	return out
}

func appendMove(transcript []Expectation,
	id string, tm int,
	move [2]*tak.Move) []Expectation {
	transcript = append(transcript, Expectation{
		recv: []string{
			fmt.Sprintf("Game#%s %s", id, playtak.FormatServer(move[0])),
		},
	})
	if move[1] == nil {
		return transcript
	}
	transcript = append(transcript, Expectation{
		send: []string{
			fmt.Sprintf("Game#%s %s", id, playtak.FormatServer(move[1])),
			fmt.Sprintf("Game#%s Time %d %d", id, tm, tm),
		},
	})
	return transcript
}

const startLine = "Game Start 100 5 Taktician vs HonestJoe white 600"

func setupGame() (*TestBotStatic, []Expectation) {
	moves := parseMoves([][2]string{
		{"a1", "e1"},
		{"e3", "b1"},
		{"e2", "b2"},
		{"Ce4", "a2"},
		{"e5", ""},
	})
	bot := &TestBotStatic{}
	for _, r := range moves {
		bot.moves = append(bot.moves, *r[0])
	}

	var transcript []Expectation
	tm := 600
	for _, r := range moves {
		transcript = appendMove(
			transcript, "100", tm, r)
		tm -= 10
	}
	transcript = append(transcript, Expectation{
		send: []string{
			"Game#100 Over R-0",
		},
	})
	return bot, transcript
}

func TestBasicGame(t *testing.T) {
	bot, transcript := setupGame()
	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}

func TestUndoGame(t *testing.T) {
	base, transcript := setupGame()
	bot := &TestBotUndo{*base, 5}

	i := 6
	rest := transcript[i:]
	transcript = transcript[:i:i]
	e := transcript[i-1]
	transcript = append(transcript,
		Expectation{
			send: []string{
				"Game#100 RequestUndo",
			},
			recv: []string{
				"Game#100 RequestUndo",
			},
		},
		Expectation{
			send: []string{
				"Game#100 Undo",
			},
		},
		Expectation{
			send: e.send,
		},
	)
	transcript = append(transcript, rest...)

	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}

func TestThinker(t *testing.T) {
	base, transcript := setupGame()
	bot := &TestBotThinker{*base}

	c := NewTestClient(t, transcript)
	playGame(c, bot, startLine)
	final := ptn.FormatTPS(bot.game.positions[len(bot.game.positions)-1])
	want := `x4,1/x4,1C/x4,1/2,2,x2,1/2,2,x2,1 2 5`
	if final != want {
		t.Fatalf("final position=%q !=%q",
			final, want)
	}
}
<|endoftext|>


<commit_before>package codegen

import (
	"fmt"
	"strconv"

	"github.com/carlcui/expressive/typing"

	"github.com/carlcui/expressive/ast"
	"github.com/carlcui/expressive/logger"
)

// CodegenVisitor visits each node and generates llvm IR.
type CodegenVisitor struct {
	logger                 logger.Logger
	labeller               *Labeller
	constants              *Fragment // global constants
	codeMap                map[ast.Node]*Fragment
	localIdentifierTracker *LocalIdentifierTracker
	externals              *Fragment
}

func (visitor *CodegenVisitor) externalFragment() *Fragment {
	fragment := NewFragment(VOID, nil)

	fragment.AddInstruction("declare i32 @printf(i8* noalias nocapture, ...) nounwind")

	return fragment
}

// Init with a logger
func (visitor *CodegenVisitor) Init(logger logger.Logger) {
	visitor.logger = logger
	visitor.labeller = &Labeller{0}
	visitor.constants = NewFragment(VOID, &GlobalIdentifierTracker{0})
	visitor.codeMap = make(map[ast.Node]*Fragment)
	visitor.localIdentifierTracker = &LocalIdentifierTracker{0}
	visitor.externals = visitor.externalFragment()
}

func (visitor *CodegenVisitor) newVoidCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(VOID, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) newValueCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(VALUE, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) newAddressCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(ADDRESS, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) getAndRemoveCode(node ast.Node) *Fragment {
	fragment, exists := visitor.codeMap[node]

	if !exists {
		panic(fmt.Sprintf("Code for node %v does not exist.", node))
	}

	delete(visitor.codeMap, node)

	return fragment
}

func (visitor *CodegenVisitor) removeVoidCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != VOID {
		panic(fmt.Sprintf("Code fragment does not produce void result: %v", node))
	}

	return fragment
}

func (visitor *CodegenVisitor) removeAddressCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != ADDRESS {
		panic(fmt.Sprintf("Code fragment does not produce address result: %v", node))
	}

	return fragment
}

func (visitor *CodegenVisitor) removeValueCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != VALUE && fragment.ResultType != ADDRESS {
		panic(fmt.Sprintf("Code fragment does not produce value result: %v", node))
	}

	if fragment.ResultType == ADDRESS {
		visitor.turnAddressIntoValue(node, fragment)
	}

	return fragment
}

func (visitor *CodegenVisitor) turnAddressIntoValue(node ast.Node, fragment *Fragment) {
	fragment.ResultType = VALUE

	result := fragment.result
	typing := node.GetTyping()
	irType := typing.IrType()

	fragment.AddOperation("load %v, %v* %v", irType, irType, result)
}

// VisitEnterProgramNode creates program scope
func (visitor *CodegenVisitor) VisitEnterProgramNode(node *ast.ProgramNode) {

}

// VisitLeaveProgramNode closes program scope
func (visitor *CodegenVisitor) VisitLeaveProgramNode(node *ast.ProgramNode) {
	// generates main function
	fragment := visitor.newVoidCode(node)

	fragment.AddInstruction("define i32 @main() {")

	for _, child := range node.Chilren {
		fragment.Append(visitor.removeVoidCode(child))
	}

	fragment.AddInstruction("ret i32 0")
	fragment.AddInstruction("}")
}

// stmts

// VisitEnterVariableDeclarationNode do something
func (visitor *CodegenVisitor) VisitEnterVariableDeclarationNode(node *ast.VariableDeclarationNode) {

}

// VisitLeaveVariableDeclarationNode do something
func (visitor *CodegenVisitor) VisitLeaveVariableDeclarationNode(node *ast.VariableDeclarationNode) {
	fragment := visitor.newVoidCode(node)

	identifierNode := node.Identifier.(*ast.IdentifierNode)
	identifierTyping := identifierNode.GetTyping()
	irType := identifierTyping.IrType()
	alignment := identifierTyping.Size()

	variable := AsLocalVariable(identifierNode.Tok.Raw)

	// allocate space
	fragment.AddInstruction("%v = alloca %v, align %v", variable, irType, alignment)

	if node.Expr == nil {
		// load default value
		defaultValue := "0"
		if identifierTyping == typing.FLOAT {
			defaultValue = "0.0"
		}

		fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, defaultValue, irType, variable, alignment)
	} else {
		exprFragment := visitor.removeValueCode(node.Expr)

		exprResultVariable := exprFragment.GetResult()

		fragment.Append(exprFragment)

		fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, exprResultVariable, irType, variable, alignment)
	}
}

// VisitEnterAssignmentNode do something
func (visitor *CodegenVisitor) VisitEnterAssignmentNode(node *ast.AssignmentNode) {
}

// VisitLeaveAssignmentNode do something
func (visitor *CodegenVisitor) VisitLeaveAssignmentNode(node *ast.AssignmentNode) {
	fragment := visitor.newVoidCode(node)

	identifierNode := node.Identifier.(*ast.IdentifierNode)
	typing := identifierNode.GetTyping()
	irType := typing.IrType()
	alignment := typing.Size()

	variable := AsLocalVariable(identifierNode.Tok.Raw)

	exprFragment := visitor.removeValueCode(node.Expr)

	exprResultVariable := exprFragment.GetResult()

	fragment.Append(exprFragment)

	fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, exprResultVariable, irType, variable, alignment)
}

// VisitEnterPrintNode do something
func (visitor *CodegenVisitor) VisitEnterPrintNode(node *ast.PrintNode) {

}

// VisitLeavePrintNode do something
func (visitor *CodegenVisitor) VisitLeavePrintNode(node *ast.PrintNode) {
	fragment := visitor.newVoidCode(node)

	stringExprFrag := visitor.removeAddressCode(node.StringExpr)

	instructionArgs := make([]interface{}, 0)

	localIdentifier := stringExprFrag.GetResult()

	instructionArgs = append(instructionArgs, localIdentifier)

	fragment.Append(stringExprFrag)

	callInstruction := "call i32 (i8*, ...) @printf(i8* %v"

	for _, arg := range node.Args {
		argFrag := visitor.removeValueCode(arg)
		localIdentifier = argFrag.GetResult()

		irType := arg.GetTyping().IrType()

		fragment.Append(argFrag)

		instructionArgs = append(instructionArgs, irType, localIdentifier)

		callInstruction += ", %v %v"
	}

	callInstruction += ")"

	fragment.AddInstruction(callInstruction, instructionArgs...)
}

// exprs

// VisitEnterTernaryOperatorNode do something
func (visitor *CodegenVisitor) VisitEnterTernaryOperatorNode(node *ast.TernaryOperatorNode) {

}

// VisitLeaveTernaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveTernaryOperatorNode(node *ast.TernaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Expr1)
	fragment2 := visitor.removeValueCode(node.Expr2)
	fragment3 := visitor.removeValueCode(node.Expr3)

	operator := node.Operator
	typing := node.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1, fragment2, fragment3)

	operatorCodegen.GenerateCode()
}

// VisitEnterBinaryOepratorNode do something
func (visitor *CodegenVisitor) VisitEnterBinaryOepratorNode(node *ast.BinaryOperatorNode) {

}

// VisitLeaveBinaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveBinaryOperatorNode(node *ast.BinaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Lhs)
	fragment2 := visitor.removeValueCode(node.Rhs)

	operator := node.Operator
	typing := node.Lhs.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1, fragment2)

	operatorCodegen.GenerateCode()
}

// VisitEnterUnaryOperatorNode do something
func (visitor *CodegenVisitor) VisitEnterUnaryOperatorNode(node *ast.UnaryOperatorNode) {

}

// VisitLeaveUnaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveUnaryOperatorNode(node *ast.UnaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Expr)

	operator := node.Operator
	typing := node.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1)

	operatorCodegen.GenerateCode()
}

// literal nodes

// VisitIntegerNode do something
func (visitor *CodegenVisitor) VisitIntegerNode(node *ast.IntegerNode) {
	fragment := visitor.newValueCode(node)

	fragment.result = strconv.Itoa(node.Val)
}

// VisitFloatNode do something
func (visitor *CodegenVisitor) VisitFloatNode(node *ast.FloatNode) {
	fragment := visitor.newValueCode(node)

	if node.Val == 0 {
		fragment.result = "0.0"
	} else {
		fragment.result = strconv.FormatFloat(float64(node.Val), 'f', 6, 32)
	}
}

// VisitCharacterNode do something
func (visitor *CodegenVisitor) VisitCharacterNode(node *ast.CharacterNode) {

}

// VisitStringNode do something
func (visitor *CodegenVisitor) VisitStringNode(node *ast.StringNode) {
	fragment := visitor.newAddressCode(node)

	stringValue := node.EscapeVal()
	stringLength := node.EscapedStringLength()

	stringGlobalIdentifier := visitor.constants.AddOperation("private constant [%v x i8] c\"%v\", align 1", stringLength, stringValue)

	fragment.AddOperation("getelementptr inbounds [%v x i8], [%v x i8]* %v, i32 0, i32 0", stringLength, stringLength, stringGlobalIdentifier)
}

// VisitIdentifierNode do something
func (visitor *CodegenVisitor) VisitIdentifierNode(node *ast.IdentifierNode) {
	fragment := visitor.newAddressCode(node)

	identifier := node.Tok.Raw

	fragment.result = AsLocalVariable(identifier)
}

// VisitBooleanNode do something
func (visitor *CodegenVisitor) VisitBooleanNode(node *ast.BooleanNode) {
	fragment := visitor.newValueCode(node)

	var booleanValue string

	if node.Val == true {
		booleanValue = "1"
	} else {
		booleanValue = "0"
	}

	fragment.result = booleanValue
}

// VisitTypeLiteralNode do something
func (visitor *CodegenVisitor) VisitTypeLiteralNode(node *ast.TypeLiteralNode) {

}

// VisitErrorNode should not happen during codegen
func (visitor *CodegenVisitor) VisitErrorNode(node *ast.ErrorNode) {
	panic(node.GetLocation() + ": unexpected error node")
}

func (visitor *CodegenVisitor) log(location string, message string) {
	visitor.logger.Log(location, message)
}
<commit_msg>add visitor funcs to codegenVisitor<commit_after>package codegen

import (
	"fmt"
	"strconv"

	"github.com/carlcui/expressive/typing"

	"github.com/carlcui/expressive/ast"
	"github.com/carlcui/expressive/logger"
)

// CodegenVisitor visits each node and generates llvm IR.
type CodegenVisitor struct {
	logger                 logger.Logger
	labeller               *Labeller
	constants              *Fragment // global constants
	codeMap                map[ast.Node]*Fragment
	localIdentifierTracker *LocalIdentifierTracker
	externals              *Fragment
}

func (visitor *CodegenVisitor) externalFragment() *Fragment {
	fragment := NewFragment(VOID, nil)

	fragment.AddInstruction("declare i32 @printf(i8* noalias nocapture, ...) nounwind")

	return fragment
}

// Init with a logger
func (visitor *CodegenVisitor) Init(logger logger.Logger) {
	visitor.logger = logger
	visitor.labeller = &Labeller{0}
	visitor.constants = NewFragment(VOID, &GlobalIdentifierTracker{0})
	visitor.codeMap = make(map[ast.Node]*Fragment)
	visitor.localIdentifierTracker = &LocalIdentifierTracker{0}
	visitor.externals = visitor.externalFragment()
}

func (visitor *CodegenVisitor) newVoidCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(VOID, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) newValueCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(VALUE, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) newAddressCode(node ast.Node) *Fragment {
	if _, exists := visitor.codeMap[node]; exists {
		panic(fmt.Sprintf("Code for node %v already exists.", node))
	}

	fragment := NewFragment(ADDRESS, visitor.localIdentifierTracker)
	visitor.codeMap[node] = fragment

	return fragment
}

func (visitor *CodegenVisitor) getAndRemoveCode(node ast.Node) *Fragment {
	fragment, exists := visitor.codeMap[node]

	if !exists {
		panic(fmt.Sprintf("Code for node %v does not exist.", node))
	}

	delete(visitor.codeMap, node)

	return fragment
}

func (visitor *CodegenVisitor) removeVoidCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != VOID {
		panic(fmt.Sprintf("Code fragment does not produce void result: %v", node))
	}

	return fragment
}

func (visitor *CodegenVisitor) removeAddressCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != ADDRESS {
		panic(fmt.Sprintf("Code fragment does not produce address result: %v", node))
	}

	return fragment
}

func (visitor *CodegenVisitor) removeValueCode(node ast.Node) *Fragment {
	fragment := visitor.getAndRemoveCode(node)

	if fragment.ResultType != VALUE && fragment.ResultType != ADDRESS {
		panic(fmt.Sprintf("Code fragment does not produce value result: %v", node))
	}

	if fragment.ResultType == ADDRESS {
		visitor.turnAddressIntoValue(node, fragment)
	}

	return fragment
}

func (visitor *CodegenVisitor) turnAddressIntoValue(node ast.Node, fragment *Fragment) {
	fragment.ResultType = VALUE

	result := fragment.result
	typing := node.GetTyping()
	irType := typing.IrType()

	fragment.AddOperation("load %v, %v* %v", irType, irType, result)
}

// VisitEnterProgramNode creates program scope
func (visitor *CodegenVisitor) VisitEnterProgramNode(node *ast.ProgramNode) {

}

// VisitLeaveProgramNode closes program scope
func (visitor *CodegenVisitor) VisitLeaveProgramNode(node *ast.ProgramNode) {
	// generates main function
	fragment := visitor.newVoidCode(node)

	fragment.AddInstruction("define i32 @main() {")

	for _, child := range node.Chilren {
		fragment.Append(visitor.removeVoidCode(child))
	}

	fragment.AddInstruction("ret i32 0")
	fragment.AddInstruction("}")
}

func (visitor *CodegenVisitor) VisitEnterBlockNode(node *ast.BlockNode) {

}

func (visitor *CodegenVisitor) VisitLeaveBlockNode(node *ast.BlockNode) {

}

// stmts

// VisitEnterVariableDeclarationNode do something
func (visitor *CodegenVisitor) VisitEnterVariableDeclarationNode(node *ast.VariableDeclarationNode) {

}

// VisitLeaveVariableDeclarationNode do something
func (visitor *CodegenVisitor) VisitLeaveVariableDeclarationNode(node *ast.VariableDeclarationNode) {
	fragment := visitor.newVoidCode(node)

	identifierNode := node.Identifier.(*ast.IdentifierNode)
	identifierTyping := identifierNode.GetTyping()
	irType := identifierTyping.IrType()
	alignment := identifierTyping.Size()

	variable := AsLocalVariable(identifierNode.Tok.Raw)

	// allocate space
	fragment.AddInstruction("%v = alloca %v, align %v", variable, irType, alignment)

	if node.Expr == nil {
		// load default value
		defaultValue := "0"
		if identifierTyping == typing.FLOAT {
			defaultValue = "0.0"
		}

		fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, defaultValue, irType, variable, alignment)
	} else {
		exprFragment := visitor.removeValueCode(node.Expr)

		exprResultVariable := exprFragment.GetResult()

		fragment.Append(exprFragment)

		fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, exprResultVariable, irType, variable, alignment)
	}
}

// VisitEnterAssignmentNode do something
func (visitor *CodegenVisitor) VisitEnterAssignmentNode(node *ast.AssignmentNode) {
}

// VisitLeaveAssignmentNode do something
func (visitor *CodegenVisitor) VisitLeaveAssignmentNode(node *ast.AssignmentNode) {
	fragment := visitor.newVoidCode(node)

	identifierNode := node.Identifier.(*ast.IdentifierNode)
	typing := identifierNode.GetTyping()
	irType := typing.IrType()
	alignment := typing.Size()

	variable := AsLocalVariable(identifierNode.Tok.Raw)

	exprFragment := visitor.removeValueCode(node.Expr)

	exprResultVariable := exprFragment.GetResult()

	fragment.Append(exprFragment)

	fragment.AddInstruction("store %v %v, %v* %v, align %v", irType, exprResultVariable, irType, variable, alignment)
}

// VisitEnterPrintNode do something
func (visitor *CodegenVisitor) VisitEnterPrintNode(node *ast.PrintNode) {

}

// VisitLeavePrintNode do something
func (visitor *CodegenVisitor) VisitLeavePrintNode(node *ast.PrintNode) {
	fragment := visitor.newVoidCode(node)

	stringExprFrag := visitor.removeAddressCode(node.StringExpr)

	instructionArgs := make([]interface{}, 0)

	localIdentifier := stringExprFrag.GetResult()

	instructionArgs = append(instructionArgs, localIdentifier)

	fragment.Append(stringExprFrag)

	callInstruction := "call i32 (i8*, ...) @printf(i8* %v"

	for _, arg := range node.Args {
		argFrag := visitor.removeValueCode(arg)
		localIdentifier = argFrag.GetResult()

		irType := arg.GetTyping().IrType()

		fragment.Append(argFrag)

		instructionArgs = append(instructionArgs, irType, localIdentifier)

		callInstruction += ", %v %v"
	}

	callInstruction += ")"

	fragment.AddInstruction(callInstruction, instructionArgs...)
}

func (visitor *CodegenVisitor) VisitEnterIfStmtNode(node *ast.IfStmtNode) {

}

func (visitor *CodegenVisitor) VisitLeaveIfStmtNode(node *ast.IfStmtNode) {

}

// exprs

// VisitEnterTernaryOperatorNode do something
func (visitor *CodegenVisitor) VisitEnterTernaryOperatorNode(node *ast.TernaryOperatorNode) {

}

// VisitLeaveTernaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveTernaryOperatorNode(node *ast.TernaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Expr1)
	fragment2 := visitor.removeValueCode(node.Expr2)
	fragment3 := visitor.removeValueCode(node.Expr3)

	operator := node.Operator
	typing := node.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1, fragment2, fragment3)

	operatorCodegen.GenerateCode()
}

// VisitEnterBinaryOepratorNode do something
func (visitor *CodegenVisitor) VisitEnterBinaryOepratorNode(node *ast.BinaryOperatorNode) {

}

// VisitLeaveBinaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveBinaryOperatorNode(node *ast.BinaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Lhs)
	fragment2 := visitor.removeValueCode(node.Rhs)

	operator := node.Operator
	typing := node.Lhs.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1, fragment2)

	operatorCodegen.GenerateCode()
}

// VisitEnterUnaryOperatorNode do something
func (visitor *CodegenVisitor) VisitEnterUnaryOperatorNode(node *ast.UnaryOperatorNode) {

}

// VisitLeaveUnaryOperatorNode do something
func (visitor *CodegenVisitor) VisitLeaveUnaryOperatorNode(node *ast.UnaryOperatorNode) {
	fragment := visitor.newValueCode(node)

	fragment1 := visitor.removeValueCode(node.Expr)

	operator := node.Operator
	typing := node.GetTyping()

	operatorCodegen := NewOperatorCodegen(fragment, operator, typing, visitor.labeller, fragment1)

	operatorCodegen.GenerateCode()
}

// literal nodes

// VisitIntegerNode do something
func (visitor *CodegenVisitor) VisitIntegerNode(node *ast.IntegerNode) {
	fragment := visitor.newValueCode(node)

	fragment.result = strconv.Itoa(node.Val)
}

// VisitFloatNode do something
func (visitor *CodegenVisitor) VisitFloatNode(node *ast.FloatNode) {
	fragment := visitor.newValueCode(node)

	if node.Val == 0 {
		fragment.result = "0.0"
	} else {
		fragment.result = strconv.FormatFloat(float64(node.Val), 'f', 6, 32)
	}
}

// VisitCharacterNode do something
func (visitor *CodegenVisitor) VisitCharacterNode(node *ast.CharacterNode) {

}

// VisitStringNode do something
func (visitor *CodegenVisitor) VisitStringNode(node *ast.StringNode) {
	fragment := visitor.newAddressCode(node)

	stringValue := node.EscapeVal()
	stringLength := node.EscapedStringLength()

	stringGlobalIdentifier := visitor.constants.AddOperation("private constant [%v x i8] c\"%v\", align 1", stringLength, stringValue)

	fragment.AddOperation("getelementptr inbounds [%v x i8], [%v x i8]* %v, i32 0, i32 0", stringLength, stringLength, stringGlobalIdentifier)
}

// VisitIdentifierNode do something
func (visitor *CodegenVisitor) VisitIdentifierNode(node *ast.IdentifierNode) {
	fragment := visitor.newAddressCode(node)

	identifier := node.Tok.Raw

	fragment.result = AsLocalVariable(identifier)
}

// VisitBooleanNode do something
func (visitor *CodegenVisitor) VisitBooleanNode(node *ast.BooleanNode) {
	fragment := visitor.newValueCode(node)

	var booleanValue string

	if node.Val == true {
		booleanValue = "1"
	} else {
		booleanValue = "0"
	}

	fragment.result = booleanValue
}

// VisitTypeLiteralNode do something
func (visitor *CodegenVisitor) VisitTypeLiteralNode(node *ast.TypeLiteralNode) {

}

// VisitErrorNode should not happen during codegen
func (visitor *CodegenVisitor) VisitErrorNode(node *ast.ErrorNode) {
	panic(node.GetLocation() + ": unexpected error node")
}

func (visitor *CodegenVisitor) log(location string, message string) {
	visitor.logger.Log(location, message)
}
<|endoftext|>


<commit_before>package stacks

import (
	"encoding/json"
	"fmt"
	"qaz/bucket"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
)

// Change - Manage Cloudformation Change-Sets
func (s *Stack) Change(req, changename string) error {
	svc := cloudformation.New(s.Session, &aws.Config{Credentials: s.creds()})

	switch req {

	case "create":
		// Resolve Deploy-Time functions
		err := s.DeployTimeParser()
		if err != nil {
			return err
		}

		params := &cloudformation.CreateChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		Log.Debug(fmt.Sprintf("Updated Template:\n%s", s.Template))

		// If bucket - upload to s3
		var (
			exists bool
			url    string
		)

		if s.Bucket != "" {
			exists, err = bucket.Exists(s.Bucket, s.Session)
			if err != nil {
				Log.Warn(fmt.Sprintf("Received Error when checking if [%s] exists: %s", s.Bucket, err.Error()))
			}
			fmt.Println("This is test")
			if !exists {
				Log.Info(fmt.Sprintf(("Creating Bucket [%s]"), s.Bucket))
				if err = bucket.Create(s.Bucket, s.Session); err != nil {
					return err
				}
			}
			t := time.Now()
			tStamp := fmt.Sprintf("%d-%d-%d_%d%d", t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute())
			url, err = bucket.S3write(s.Bucket, fmt.Sprintf("%s_%s.template", s.Stackname, tStamp), s.Template, s.Session)
			if err != nil {
				return err
			}
			params.TemplateURL = &url
		} else {
			params.TemplateBody = &s.Template
		}

		// If IAM is bening touched, add Capabilities
		if strings.Contains(s.Template, "AWS::IAM") {
			params.Capabilities = []*string{
				aws.String(cloudformation.CapabilityCapabilityIam),
				aws.String(cloudformation.CapabilityCapabilityNamedIam),
			}
		}

		if _, err = svc.CreateChangeSet(params); err != nil {
			return err
		}

		describeParams := &cloudformation.DescribeChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		for {
			// Waiting for PENDING state to change
			resp, err := svc.DescribeChangeSet(describeParams)
			if err != nil {
				return err
			}

			Log.Info(fmt.Sprintf("Creating Change-Set: [%s] - %s - %s", changename, Log.ColorMap(*resp.Status), s.Stackname))

			if *resp.Status == "CREATE_COMPLETE" || *resp.Status == "FAILED" {
				break
			}

			time.Sleep(time.Second * 1)
		}

	case "rm":
		params := &cloudformation.DeleteChangeSetInput{
			ChangeSetName: aws.String(changename),
			StackName:     aws.String(s.Stackname),
		}

		if _, err := svc.DeleteChangeSet(params); err != nil {
			return err
		}

		Log.Info(fmt.Sprintf("Change-Set: [%s] deleted", changename))

	case "list":
		params := &cloudformation.ListChangeSetsInput{
			StackName: aws.String(s.Stackname),
		}

		resp, err := svc.ListChangeSets(params)
		if err != nil {
			return err
		}

		for _, i := range resp.Summaries {
			Log.Info(fmt.Sprintf("%s%s - Change-Set: [%s] - Status: [%s]", Log.ColorString("@", "magenta"), i.CreationTime.Format(time.RFC850), *i.ChangeSetName, *i.ExecutionStatus))
		}

	case "execute":
		done := make(chan bool)
		params := &cloudformation.ExecuteChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		if _, err := svc.ExecuteChangeSet(params); err != nil {
			return err
		}

		describeStacksInput := &cloudformation.DescribeStacksInput{
			StackName: aws.String(s.Stackname),
		}

		go s.tail("UPDATE", done)

		Log.Debug(fmt.Sprintln("Calling [WaitUntilStackUpdateComplete] with parameters:", describeStacksInput))
		if err := svc.WaitUntilStackUpdateComplete(describeStacksInput); err != nil {
			return err
		}

		done <- true

	case "desc":
		params := &cloudformation.DescribeChangeSetInput{
			ChangeSetName: aws.String(changename),
			StackName:     aws.String(s.Stackname),
		}

		resp, err := svc.DescribeChangeSet(params)
		if err != nil {
			return err
		}

		o, err := json.MarshalIndent(resp, "", "  ")
		if err != nil {
			return err
		}

		fmt.Printf("%s\n", o)
	}

	return nil
}
<commit_msg>added parameters to change-set create<commit_after>package stacks

import (
	"encoding/json"
	"fmt"
	"qaz/bucket"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
)

// Change - Manage Cloudformation Change-Sets
func (s *Stack) Change(req, changename string) error {
	svc := cloudformation.New(s.Session, &aws.Config{Credentials: s.creds()})

	switch req {

	case "create":
		// Resolve Deploy-Time functions
		err := s.DeployTimeParser()
		if err != nil {
			return err
		}

		params := &cloudformation.CreateChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		// add tags if set
		if len(s.Tags) > 0 {
			params.Tags = s.Tags
		}

		Log.Debug(fmt.Sprintf("Updated Template:\n%s", s.Template))

		// If bucket - upload to s3
		var (
			exists bool
			url    string
		)

		if s.Bucket != "" {
			exists, err = bucket.Exists(s.Bucket, s.Session)
			if err != nil {
				Log.Warn(fmt.Sprintf("Received Error when checking if [%s] exists: %s", s.Bucket, err.Error()))
			}
			fmt.Println("This is test")
			if !exists {
				Log.Info(fmt.Sprintf(("Creating Bucket [%s]"), s.Bucket))
				if err = bucket.Create(s.Bucket, s.Session); err != nil {
					return err
				}
			}
			t := time.Now()
			tStamp := fmt.Sprintf("%d-%d-%d_%d%d", t.Year(), t.Month(), t.Day(), t.Hour(), t.Minute())
			url, err = bucket.S3write(s.Bucket, fmt.Sprintf("%s_%s.template", s.Stackname, tStamp), s.Template, s.Session)
			if err != nil {
				return err
			}
			params.TemplateURL = &url
		} else {
			params.TemplateBody = &s.Template
		}

		// If IAM is bening touched, add Capabilities
		if strings.Contains(s.Template, "AWS::IAM") {
			params.Capabilities = []*string{
				aws.String(cloudformation.CapabilityCapabilityIam),
				aws.String(cloudformation.CapabilityCapabilityNamedIam),
			}
		}

		if _, err = svc.CreateChangeSet(params); err != nil {
			return err
		}

		describeParams := &cloudformation.DescribeChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		for {
			// Waiting for PENDING state to change
			resp, err := svc.DescribeChangeSet(describeParams)
			if err != nil {
				return err
			}

			Log.Info(fmt.Sprintf("Creating Change-Set: [%s] - %s - %s", changename, Log.ColorMap(*resp.Status), s.Stackname))

			if *resp.Status == "CREATE_COMPLETE" || *resp.Status == "FAILED" {
				break
			}

			time.Sleep(time.Second * 1)
		}

	case "rm":
		params := &cloudformation.DeleteChangeSetInput{
			ChangeSetName: aws.String(changename),
			StackName:     aws.String(s.Stackname),
		}

		if _, err := svc.DeleteChangeSet(params); err != nil {
			return err
		}

		Log.Info(fmt.Sprintf("Change-Set: [%s] deleted", changename))

	case "list":
		params := &cloudformation.ListChangeSetsInput{
			StackName: aws.String(s.Stackname),
		}

		resp, err := svc.ListChangeSets(params)
		if err != nil {
			return err
		}

		for _, i := range resp.Summaries {
			Log.Info(fmt.Sprintf("%s%s - Change-Set: [%s] - Status: [%s]", Log.ColorString("@", "magenta"), i.CreationTime.Format(time.RFC850), *i.ChangeSetName, *i.ExecutionStatus))
		}

	case "execute":
		done := make(chan bool)
		params := &cloudformation.ExecuteChangeSetInput{
			StackName:     aws.String(s.Stackname),
			ChangeSetName: aws.String(changename),
		}

		if _, err := svc.ExecuteChangeSet(params); err != nil {
			return err
		}

		describeStacksInput := &cloudformation.DescribeStacksInput{
			StackName: aws.String(s.Stackname),
		}

		go s.tail("UPDATE", done)

		Log.Debug(fmt.Sprintln("Calling [WaitUntilStackUpdateComplete] with parameters:", describeStacksInput))
		if err := svc.WaitUntilStackUpdateComplete(describeStacksInput); err != nil {
			return err
		}

		done <- true

	case "desc":
		params := &cloudformation.DescribeChangeSetInput{
			ChangeSetName: aws.String(changename),
			StackName:     aws.String(s.Stackname),
		}

		resp, err := svc.DescribeChangeSet(params)
		if err != nil {
			return err
		}

		o, err := json.MarshalIndent(resp, "", "  ")
		if err != nil {
			return err
		}

		fmt.Printf("%s\n", o)
	}

	return nil
}
<|endoftext|>


<commit_before>package lastfm

import (
	"fmt"
	"github.com/0x263b/Porygon2"
	"github.com/0x263b/Porygon2/web"
	"strings"
)

const (
	ChartsURL = "http://ws.audioscrobbler.com/2.0/?method=user.gettopartists&user=%s&period=7day&limit=5&api_key=%s&format=json"
)

func charts(command *bot.Cmd, matches []string) (msg string, err error) {
	username := checkLastfm(command.Nick, matches[1])

	if username == "" {
		return "Lastfm not provided, nor on file. Use `-set lastfm <lastfm>` to save", nil
	}

	data := &WeeklyCharts{}
	err = web.GetJSON(fmt.Sprintf(ChartsURL, username, bot.Config.API.Lastfm), data)
	if err != nil || data.Error > 0 {
		return fmt.Sprintf("Could not get charts for %s", username), nil
	}
	if data.Topartists.Attr.Total == "0" {
		return fmt.Sprintf("Could not get charts for %s", username), nil
	}

	var fmtcharts string
	for i := range data.Topartists.Artist[:5] {
		fmtcharts += fmt.Sprintf("%s (%s), ", data.Topartists.Artist[i].Name, data.Topartists.Artist[i].Playcount)
	}
	fmtcharts = strings.TrimSuffix(fmtcharts, ", ")

	output := fmt.Sprintf("Last.fm | Top 5 Weekly artists for %s | %s", username, fmtcharts)

	return output, nil
}

func init() {
	bot.RegisterCommand(
		"^charts(?: (\\S+))?$",
		charts)
}
<commit_msg>Shallow charts<commit_after>package lastfm

import (
	"fmt"
	"github.com/0x263b/Porygon2"
	"github.com/0x263b/Porygon2/web"
	"strings"
)

const (
	ChartsURL = "http://ws.audioscrobbler.com/2.0/?method=user.gettopartists&user=%s&period=7day&limit=5&api_key=%s&format=json"
)

func charts(command *bot.Cmd, matches []string) (msg string, err error) {
	username := checkLastfm(command.Nick, matches[1])

	if username == "" {
		return "Lastfm not provided, nor on file. Use `-set lastfm <lastfm>` to save", nil
	}

	data := &WeeklyCharts{}
	err = web.GetJSON(fmt.Sprintf(ChartsURL, username, bot.Config.API.Lastfm), data)
	if err != nil || data.Error > 0 {
		return fmt.Sprintf("Could not get charts for %s", username), nil
	}
	if data.Topartists.Attr.Total == "0" {
		return fmt.Sprintf("Could not get charts for %s", username), nil
	}

	var fmtcharts string
	var trunc int = 5

	if len(data.Topartists.Artist) < trunc {
		trunc = len(data.Topartists.Artist)
	}

	for i := range data.Topartists.Artist[:trunc] {
		fmtcharts += fmt.Sprintf("%s (%s), ", data.Topartists.Artist[i].Name, data.Topartists.Artist[i].Playcount)
	}
	fmtcharts = strings.TrimSuffix(fmtcharts, ", ")

	output := fmt.Sprintf("Last.fm | Top 5 Weekly artists for %s | %s", username, fmtcharts)

	return output, nil
}

func init() {
	bot.RegisterCommand(
		"^charts(?: (\\S+))?$",
		charts)
}
<|endoftext|>