perf tools: Add sched map changes

index 267112b4e3db..878ac0687b0a 100644
@@ -436,3 +436,15 @@ struct thread_map *thread_map__new_event(struct thread_map_event *event)
 
 	return threads;
 }
+
+bool thread_map__has(struct thread_map *threads, pid_t pid)
+{
+	int i;
+
+	for (i = 0; i < threads->nr; ++i) {
+		if (threads->map[i].pid == pid)
+			return true;
+	}
+
+	return false;
+}

index 85e4c7c4fbde..9a065ea69ff1 100644
@@ -55,4 +55,5 @@ static inline char *thread_map__comm(struct thread_map *map, int thread)
 }
 
 void thread_map__read_comms(struct thread_map *threads);
+bool thread_map__has(struct thread_map *threads, pid_t pid);
 #endif	/* __PERF_THREAD_MAP_H */

index 9bcf2bed3a6d..02d801670f30 100644
@@ -587,3 +587,15 @@ int cpu__setup_cpunode_map(void)
 	closedir(dir1);
 	return 0;
 }
+
+bool cpu_map__has(struct cpu_map *cpus, int cpu)
+{
+	int i;
+
+	for (i = 0; i < cpus->nr; ++i) {
+		if (cpus->map[i] == cpu)
+			return true;
+	}
+
+	return false;
+}

index 81a2562aaa2b..1a0a35073ce1 100644
@@ -66,4 +66,6 @@ int cpu__get_node(int cpu);
 int cpu_map__build_map(struct cpu_map *cpus, struct cpu_map **res,
 		       int (*f)(struct cpu_map *map, int cpu, void *data),
 		       void *data);
+
+bool cpu_map__has(struct cpu_map *cpus, int cpu);
 #endif /* __PERF_CPUMAP_H */

index 871b55ae22a4..64dd94667055 100644
@@ -122,6 +122,12 @@ struct trace_sched_handler {
 				  struct machine *machine);
 };
 
+struct perf_sched_map {
+	DECLARE_BITMAP(comp_cpus_mask, MAX_CPUS);
+	int			*comp_cpus;
+	bool			 comp;
+};
+
 struct perf_sched {
 	struct perf_tool tool;
 	const char	 *sort_order;

@@ -173,6 +179,7 @@ struct perf_sched {
 	struct list_head sort_list, cmp_pid;
 	bool force;
 	bool skip_merge;
+	struct perf_sched_map map;
 };
 
 static u64 get_nsecs(void)

@@ -1347,13 +1354,24 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 	int new_shortname;
 	u64 timestamp0, timestamp = sample->time;
 	s64 delta;
-	int cpu, this_cpu = sample->cpu;
+	int i, this_cpu = sample->cpu;
+	int cpus_nr;
+	bool new_cpu = false;
 
 	BUG_ON(this_cpu >= MAX_CPUS || this_cpu < 0);
 
 	if (this_cpu > sched->max_cpu)
 		sched->max_cpu = this_cpu;
 
+	if (sched->map.comp) {
+		cpus_nr = bitmap_weight(sched->map.comp_cpus_mask, MAX_CPUS);
+		if (!test_and_set_bit(this_cpu, sched->map.comp_cpus_mask)) {
+			sched->map.comp_cpus[cpus_nr++] = this_cpu;
+			new_cpu = true;
+		}
+	} else
+		cpus_nr = sched->max_cpu;
+
 	timestamp0 = sched->cpu_last_switched[this_cpu];
 	sched->cpu_last_switched[this_cpu] = timestamp;
 	if (timestamp0)

@@ -1400,7 +1418,9 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 		new_shortname = 1;
 	}
 
-	for (cpu = 0; cpu <= sched->max_cpu; cpu++) {
+	for (i = 0; i < cpus_nr; i++) {
+		int cpu = sched->map.comp ? sched->map.comp_cpus[i] : i;
+
 		if (cpu != this_cpu)
 			printf(" ");
 		else

@@ -1414,12 +1434,15 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 
 	printf("  %12.6f secs ", (double)timestamp/1e9);
 	if (new_shortname) {
-		printf("%s => %s:%d\n",
+		printf("%s => %s:%d",
 		       sched_in->shortname, thread__comm_str(sched_in), sched_in->tid);
-	} else {
-		printf("\n");
 	}
 
+	if (sched->map.comp && new_cpu)
+		printf(" (CPU %d)", this_cpu);
+
+	printf("\n");
+
 	thread__put(sched_in);
 
 	return 0;

@@ -1675,9 +1698,22 @@ static int perf_sched__lat(struct perf_sched *sched)
 	return 0;
 }
 
+static int setup_map_cpus(struct perf_sched *sched)
+{
+	sched->max_cpu  = sysconf(_SC_NPROCESSORS_CONF);
+
+	if (sched->map.comp) {
+		sched->map.comp_cpus = zalloc(sched->max_cpu * sizeof(int));
+		return sched->map.comp_cpus ? 0 : -1;
+	}
+
+	return 0;
+}
+
 static int perf_sched__map(struct perf_sched *sched)
 {
-	sched->max_cpu = sysconf(_SC_NPROCESSORS_CONF);
+	if (setup_map_cpus(sched))
+		return -1;
 
 	setup_pager();
 	if (perf_sched__read_events(sched))

@@ -1831,6 +1867,11 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		    "dump raw trace in ASCII"),
 	OPT_END()
 	};
+	const struct option map_options[] = {
+	OPT_BOOLEAN(0, "compact", &sched.map.comp,
+		    "map output in compact mode"),
+	OPT_END()
+	};
 	const char * const latency_usage[] = {
 		"perf sched latency [<options>]",
 		NULL

@@ -1839,6 +1880,10 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		"perf sched replay [<options>]",
 		NULL
 	};
+	const char * const map_usage[] = {
+		"perf sched map [<options>]",
+		NULL
+	};
 	const char *const sched_subcommands[] = { "record", "latency", "map",
 						  "replay", "script", NULL };
 	const char *sched_usage[] = {

@@ -1887,6 +1932,11 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		setup_sorting(&sched, latency_options, latency_usage);
 		return perf_sched__lat(&sched);
 	} else if (!strcmp(argv[0], "map")) {
+		if (argc) {
+			argc = parse_options(argc, argv, map_options, replay_usage, 0);
+			if (argc)
+				usage_with_options(map_usage, map_options);
+		}
 		sched.tp_handler = &map_ops;
 		setup_sorting(&sched, latency_options, latency_usage);
 		return perf_sched__map(&sched);

index 64dd94667055..9ef28973f198 100644
@@ -11,6 +11,7 @@
 #include "util/session.h"
 #include "util/tool.h"
 #include "util/cloexec.h"
+#include "util/color.h"
 
 #include <subcmd/parse-options.h>
 #include "util/trace-event.h"

@@ -1357,6 +1358,7 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 	int i, this_cpu = sample->cpu;
 	int cpus_nr;
 	bool new_cpu = false;
+	const char *color = PERF_COLOR_NORMAL;
 
 	BUG_ON(this_cpu >= MAX_CPUS || this_cpu < 0);
 

@@ -1422,26 +1424,26 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 		int cpu = sched->map.comp ? sched->map.comp_cpus[i] : i;
 
 		if (cpu != this_cpu)
-			printf(" ");
+			color_fprintf(stdout, color, " ");
 		else
-			printf("*");
+			color_fprintf(stdout, color, "*");
 
 		if (sched->curr_thread[cpu])
-			printf("%2s ", sched->curr_thread[cpu]->shortname);
+			color_fprintf(stdout, color, "%2s ", sched->curr_thread[cpu]->shortname);
 		else
-			printf("   ");
+			color_fprintf(stdout, color, "   ");
 	}
 
-	printf("  %12.6f secs ", (double)timestamp/1e9);
+	color_fprintf(stdout, color, "  %12.6f secs ", (double)timestamp/1e9);
 	if (new_shortname) {
-		printf("%s => %s:%d",
+		color_fprintf(stdout, color, "%s => %s:%d",
 		       sched_in->shortname, thread__comm_str(sched_in), sched_in->tid);
 	}
 
 	if (sched->map.comp && new_cpu)
-		printf(" (CPU %d)", this_cpu);
+		color_fprintf(stdout, color, " (CPU %d)", this_cpu);
 
-	printf("\n");
+	color_fprintf(stdout, color, "\n");
 
 	thread__put(sched_in);
 

index 878ac0687b0a..5654fe15e036 100644
@@ -260,7 +260,7 @@ struct thread_map *thread_map__new_dummy(void)
 	return threads;
 }
 
-static struct thread_map *thread_map__new_by_tid_str(const char *tid_str)
+struct thread_map *thread_map__new_by_tid_str(const char *tid_str)
 {
 	struct thread_map *threads = NULL, *nt;
 	int ntasks = 0;

index 9a065ea69ff1..bd3b971588da 100644
@@ -31,6 +31,8 @@ void thread_map__put(struct thread_map *map);
 struct thread_map *thread_map__new_str(const char *pid,
 		const char *tid, uid_t uid);
 
+struct thread_map *thread_map__new_by_tid_str(const char *tid_str);
+
 size_t thread_map__fprintf(struct thread_map *threads, FILE *fp);
 
 static inline int thread_map__nr(struct thread_map *threads)

index 9ef28973f198..b5361a1d20e1 100644
@@ -11,6 +11,7 @@
 #include "util/session.h"
 #include "util/tool.h"
 #include "util/cloexec.h"
+#include "util/thread_map.h"
 #include "util/color.h"
 
 #include <subcmd/parse-options.h>

@@ -123,10 +124,14 @@ struct trace_sched_handler {
 				  struct machine *machine);
 };
 
+#define COLOR_PIDS PERF_COLOR_BLUE
+
 struct perf_sched_map {
 	DECLARE_BITMAP(comp_cpus_mask, MAX_CPUS);
 	int			*comp_cpus;
 	bool			 comp;
+	struct thread_map	*color_pids;
+	const char		*color_pids_str;
 };
 
 struct perf_sched {

@@ -1347,6 +1352,38 @@ static int process_sched_wakeup_event(struct perf_tool *tool,
 	return 0;
 }
 
+union map_priv {
+	void	*ptr;
+	bool	 color;
+};
+
+static bool thread__has_color(struct thread *thread)
+{
+	union map_priv priv = {
+		.ptr = thread__priv(thread),
+	};
+
+	return priv.color;
+}
+
+static struct thread*
+map__findnew_thread(struct perf_sched *sched, struct machine *machine, pid_t pid, pid_t tid)
+{
+	struct thread *thread = machine__findnew_thread(machine, pid, tid);
+	union map_priv priv = {
+		.color = false,
+	};
+
+	if (!sched->map.color_pids || !thread || thread__priv(thread))
+		return thread;
+
+	if (thread_map__has(sched->map.color_pids, tid))
+		priv.color = true;
+
+	thread__set_priv(thread, priv.ptr);
+	return thread;
+}
+
 static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 			    struct perf_sample *sample, struct machine *machine)
 {

@@ -1386,7 +1423,7 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 		return -1;
 	}
 
-	sched_in = machine__findnew_thread(machine, -1, next_pid);
+	sched_in = map__findnew_thread(sched, machine, -1, next_pid);
 	if (sched_in == NULL)
 		return -1;
 

@@ -1422,6 +1459,11 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 
 	for (i = 0; i < cpus_nr; i++) {
 		int cpu = sched->map.comp ? sched->map.comp_cpus[i] : i;
+		struct thread *curr_thread = sched->curr_thread[cpu];
+		const char *pid_color = color;
+
+		if (curr_thread && thread__has_color(curr_thread))
+			pid_color = COLOR_PIDS;
 
 		if (cpu != this_cpu)
 			color_fprintf(stdout, color, " ");

@@ -1429,14 +1471,19 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 			color_fprintf(stdout, color, "*");
 
 		if (sched->curr_thread[cpu])
-			color_fprintf(stdout, color, "%2s ", sched->curr_thread[cpu]->shortname);
+			color_fprintf(stdout, pid_color, "%2s ", sched->curr_thread[cpu]->shortname);
 		else
 			color_fprintf(stdout, color, "   ");
 	}
 
 	color_fprintf(stdout, color, "  %12.6f secs ", (double)timestamp/1e9);
 	if (new_shortname) {
-		color_fprintf(stdout, color, "%s => %s:%d",
+		const char *pid_color = color;
+
+		if (thread__has_color(sched_in))
+			pid_color = COLOR_PIDS;
+
+		color_fprintf(stdout, pid_color, "%s => %s:%d",
 		       sched_in->shortname, thread__comm_str(sched_in), sched_in->tid);
 	}
 

@@ -1712,11 +1759,31 @@ static int setup_map_cpus(struct perf_sched *sched)
 	return 0;
 }
 
+static int setup_color_pids(struct perf_sched *sched)
+{
+	struct thread_map *map;
+
+	if (!sched->map.color_pids_str)
+		return 0;
+
+	map = thread_map__new_by_tid_str(sched->map.color_pids_str);
+	if (!map) {
+		pr_err("failed to get thread map from %s\n", sched->map.color_pids_str);
+		return -1;
+	}
+
+	sched->map.color_pids = map;
+	return 0;
+}
+
 static int perf_sched__map(struct perf_sched *sched)
 {
 	if (setup_map_cpus(sched))
 		return -1;
 
+	if (setup_color_pids(sched))
+		return -1;
+
 	setup_pager();
 	if (perf_sched__read_events(sched))
 		return -1;

@@ -1872,6 +1939,8 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 	const struct option map_options[] = {
 	OPT_BOOLEAN(0, "compact", &sched.map.comp,
 		    "map output in compact mode"),
+	OPT_STRING(0, "color-pids", &sched.map.color_pids_str, "pids",
+		   "highlight given pids in map"),
 	OPT_END()
 	};
 	const char * const latency_usage[] = {

@@ -1935,7 +2004,7 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		return perf_sched__lat(&sched);
 	} else if (!strcmp(argv[0], "map")) {
 		if (argc) {
-			argc = parse_options(argc, argv, map_options, replay_usage, 0);
+			argc = parse_options(argc, argv, map_options, map_usage, 0);
 			if (argc)
 				usage_with_options(map_usage, map_options);
 		}

index b5361a1d20e1..7de04b297c14 100644
@@ -125,6 +125,7 @@ struct trace_sched_handler {
 };
 
 #define COLOR_PIDS PERF_COLOR_BLUE
+#define COLOR_CPUS PERF_COLOR_BG_RED
 
 struct perf_sched_map {
 	DECLARE_BITMAP(comp_cpus_mask, MAX_CPUS);

@@ -132,6 +133,8 @@ struct perf_sched_map {
 	bool			 comp;
 	struct thread_map	*color_pids;
 	const char		*color_pids_str;
+	struct cpu_map		*color_cpus;
+	const char		*color_cpus_str;
 };
 
 struct perf_sched {

@@ -1461,14 +1464,18 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 		int cpu = sched->map.comp ? sched->map.comp_cpus[i] : i;
 		struct thread *curr_thread = sched->curr_thread[cpu];
 		const char *pid_color = color;
+		const char *cpu_color = color;
 
 		if (curr_thread && thread__has_color(curr_thread))
 			pid_color = COLOR_PIDS;
 
+		if (sched->map.color_cpus && cpu_map__has(sched->map.color_cpus, cpu))
+			cpu_color = COLOR_CPUS;
+
 		if (cpu != this_cpu)
-			color_fprintf(stdout, color, " ");
+			color_fprintf(stdout, cpu_color, " ");
 		else
-			color_fprintf(stdout, color, "*");
+			color_fprintf(stdout, cpu_color, "*");
 
 		if (sched->curr_thread[cpu])
 			color_fprintf(stdout, pid_color, "%2s ", sched->curr_thread[cpu]->shortname);

@@ -1753,7 +1760,8 @@ static int setup_map_cpus(struct perf_sched *sched)
 
 	if (sched->map.comp) {
 		sched->map.comp_cpus = zalloc(sched->max_cpu * sizeof(int));
-		return sched->map.comp_cpus ? 0 : -1;
+		if (!sched->map.comp_cpus)
+			return -1;
 	}
 
 	return 0;

@@ -1776,6 +1784,23 @@ static int setup_color_pids(struct perf_sched *sched)
 	return 0;
 }
 
+static int setup_color_cpus(struct perf_sched *sched)
+{
+	struct cpu_map *map;
+
+	if (!sched->map.color_cpus_str)
+		return 0;
+
+	map = cpu_map__new(sched->map.color_cpus_str);
+	if (!map) {
+		pr_err("failed to get thread map from %s\n", sched->map.color_cpus_str);
+		return -1;
+	}
+
+	sched->map.color_cpus = map;
+	return 0;
+}
+
 static int perf_sched__map(struct perf_sched *sched)
 {
 	if (setup_map_cpus(sched))

@@ -1784,6 +1809,9 @@ static int perf_sched__map(struct perf_sched *sched)
 	if (setup_color_pids(sched))
 		return -1;
 
+	if (setup_color_cpus(sched))
+		return -1;
+
 	setup_pager();
 	if (perf_sched__read_events(sched))
 		return -1;

@@ -1941,6 +1969,8 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		    "map output in compact mode"),
 	OPT_STRING(0, "color-pids", &sched.map.color_pids_str, "pids",
 		   "highlight given pids in map"),
+	OPT_STRING(0, "color-cpus", &sched.map.color_cpus_str, "cpus",
+                    "highlight given CPUs in map"),
 	OPT_END()
 	};
 	const char * const latency_usage[] = {

index 7de04b297c14..afa057666c2a 100644
@@ -135,6 +135,8 @@ struct perf_sched_map {
 	const char		*color_pids_str;
 	struct cpu_map		*color_cpus;
 	const char		*color_cpus_str;
+	struct cpu_map		*cpus;
+	const char		*cpus_str;
 };
 
 struct perf_sched {

@@ -1469,6 +1471,9 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 		if (curr_thread && thread__has_color(curr_thread))
 			pid_color = COLOR_PIDS;
 
+		if (sched->map.cpus && !cpu_map__has(sched->map.cpus, cpu))
+			continue;
+
 		if (sched->map.color_cpus && cpu_map__has(sched->map.color_cpus, cpu))
 			cpu_color = COLOR_CPUS;
 

@@ -1483,6 +1488,9 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 			color_fprintf(stdout, color, "   ");
 	}
 
+	if (sched->map.cpus && !cpu_map__has(sched->map.cpus, this_cpu))
+		goto out;
+
 	color_fprintf(stdout, color, "  %12.6f secs ", (double)timestamp/1e9);
 	if (new_shortname) {
 		const char *pid_color = color;

@@ -1497,6 +1505,7 @@ static int map_switch_event(struct perf_sched *sched, struct perf_evsel *evsel,
 	if (sched->map.comp && new_cpu)
 		color_fprintf(stdout, color, " (CPU %d)", this_cpu);
 
+out:
 	color_fprintf(stdout, color, "\n");
 
 	thread__put(sched_in);

@@ -1756,6 +1765,8 @@ static int perf_sched__lat(struct perf_sched *sched)
 
 static int setup_map_cpus(struct perf_sched *sched)
 {
+	struct cpu_map *map;
+
 	sched->max_cpu  = sysconf(_SC_NPROCESSORS_CONF);
 
 	if (sched->map.comp) {

@@ -1764,6 +1775,16 @@ static int setup_map_cpus(struct perf_sched *sched)
 			return -1;
 	}
 
+	if (!sched->map.cpus_str)
+		return 0;
+
+	map = cpu_map__new(sched->map.cpus_str);
+	if (!map) {
+		pr_err("failed to get cpus map from %s\n", sched->map.cpus_str);
+		return -1;
+	}
+
+	sched->map.cpus = map;
 	return 0;
 }
 

@@ -1971,6 +1992,8 @@ int cmd_sched(int argc, const char **argv, const char *prefix __maybe_unused)
 		   "highlight given pids in map"),
 	OPT_STRING(0, "color-cpus", &sched.map.color_cpus_str, "cpus",
                     "highlight given CPUs in map"),
+	OPT_STRING(0, "cpus", &sched.map.cpus_str, "cpus",
+                    "display given CPUs in map"),
 	OPT_END()
 	};
 	const char * const latency_usage[] = {