#!/usr/bin/env perl
# SPDX-License-Identifier: GPL-2.0
#
# Generates a linker script that specifies the correct initcall order.
#
# Copyright (C) 2019 Google LLC

use strict;
use warnings;
use IO::Handle;
use IO::Select;
use POSIX ":sys_wait_h";

my $nm = $ENV{'NM'} || die "$0: ERROR: NM not set?";
my $objtree = $ENV{'objtree'} || '.';

## currently active child processes
my $jobs = {};		# child process pid -> file handle
## results from child processes
my $results = {};	# object index -> [ { level, secname }, ... ]

## reads _NPROCESSORS_ONLN to determine the maximum number of processes to
## start
sub get_online_processors {
	open(my $fh, "getconf _NPROCESSORS_ONLN 2>/dev/null |")
		or die "$0: ERROR: failed to execute getconf: $!";
	my $procs = <$fh>;
	close($fh);

	if (!($procs =~ /^\d+$/)) {
		return 1;
	}

	return int($procs);
}

## writes results to the parent process
## format: <file index> <initcall level> <base initcall section name>
sub write_results {
	my ($index, $initcalls) = @_;

	# sort by the counter value to ensure the order of initcalls within
	# each object file is correct
	foreach my $counter (sort { $a <=> $b } keys(%{$initcalls})) {
		my $level = $initcalls->{$counter}->{'level'};

		# section name for the initcall function
		my $secname = $initcalls->{$counter}->{'module'} . '__' .
			      $counter . '_' .
			      $initcalls->{$counter}->{'line'} . '_' .
			      $initcalls->{$counter}->{'function'};

		print "$index $level $secname\n";
	}
}

## reads a result line from a child process and adds it to the $results array
sub read_results{
	my ($fh) = @_;

	# each child prints out a full line w/ autoflush and exits after the
	# last line, so even if buffered I/O blocks here, it shouldn't block
	# very long
	my $data = <$fh>;

	if (!defined($data)) {
		return 0;
	}

	chomp($data);

	my ($index, $level, $secname) = $data =~
		/^(\d+)\ ([^\ ]+)\ (.*)$/;

	if (!defined($index) ||
		!defined($level) ||
		!defined($secname)) {
		die "$0: ERROR: child process returned invalid data: $data\n";
	}

	$index = int($index);

	if (!exists($results->{$index})) {
		$results->{$index} = [];
	}

	push (@{$results->{$index}}, {
		'level'   => $level,
		'secname' => $secname
	});

	return 1;
}

## finds initcalls from an object file or all object files in an archive, and
## writes results back to the parent process
sub find_initcalls {
	my ($index, $file) = @_;

	die "$0: ERROR: file $file doesn't exist?" if (! -f $file);

	open(my $fh, "\"$nm\" --defined-only \"$file\" 2>/dev/null |")
		or die "$0: ERROR: failed to execute \"$nm\": $!";

	my $initcalls = {};

	while (<$fh>) {
		chomp;

		# check for the start of a new object file (if processing an
		# archive)
		my ($path)= $_ =~ /^(.+)\:$/;

		if (defined($path)) {
			write_results($index, $initcalls);
			$initcalls = {};
			next;
		}

		# look for an initcall
		my ($module, $counter, $line, $symbol) = $_ =~
			/[a-z]\s+__initcall__(\S*)__(\d+)_(\d+)_(.*)$/;

		if (!defined($module)) {
			$module = ''
		}

		if (!defined($counter) ||
			!defined($line) ||
			!defined($symbol)) {
			next;
		}

		# parse initcall level
		my ($function, $level) = $symbol =~
			/^(.*)((early|rootfs|con|[0-9])s?)$/;

		die "$0: ERROR: invalid initcall name $symbol in $file($path)"
			if (!defined($function) || !defined($level));

		$initcalls->{$counter} = {
			'module'   => $module,
			'line'     => $line,
			'function' => $function,
			'level'    => $level,
		};
	}

	close($fh);
	write_results($index, $initcalls);
}

## waits for any child process to complete, reads the results, and adds them to
## the $results array for later processing
sub wait_for_results {
	my ($select) = @_;

	my $pid = 0;
	do {
		# unblock children that may have a full write buffer
		foreach my $fh ($select->can_read(0)) {
			read_results($fh);
		}

		# check for children that have exited, read the remaining data
		# from them, and clean up
		$pid = waitpid(-1, WNOHANG);
		if ($pid > 0) {
			if (!exists($jobs->{$pid})) {
				next;
			}

			my $fh = $jobs->{$pid};
			$select->remove($fh);

			while (read_results($fh)) {
				# until eof
			}

			close($fh);
			delete($jobs->{$pid});
		}
	} while ($pid > 0);
}

## forks a child to process each file passed in the command line and collects
## the results
sub process_files {
	my $index = 0;
	my $njobs = $ENV{'PARALLELISM'} || get_online_processors();
	my $select = IO::Select->new();

	while (my $file = shift(@ARGV)) {
		# fork a child process and read it's stdout
		my $pid = open(my $fh, '-|');

		if (!defined($pid)) {
			die "$0: ERROR: failed to fork: $!";
		} elsif ($pid) {
			# save the child process pid and the file handle
			$select->add($fh);
			$jobs->{$pid} = $fh;
		} else {
			# in the child process
			STDOUT->autoflush(1);
			find_initcalls($index, "$objtree/$file");
			exit;
		}

		$index++;

		# limit the number of children to $njobs
		if (scalar(keys(%{$jobs})) >= $njobs) {
			wait_for_results($select);
		}
	}

	# wait for the remaining children to complete
	while (scalar(keys(%{$jobs})) > 0) {
		wait_for_results($select);
	}
}

sub generate_initcall_lds() {
	process_files();

	my $sections = {};	# level -> [ secname, ...]

	# sort results to retain link order and split to sections per
	# initcall level
	foreach my $index (sort { $a <=> $b } keys(%{$results})) {
		foreach my $result (@{$results->{$index}}) {
			my $level = $result->{'level'};

			if (!exists($sections->{$level})) {
				$sections->{$level} = [];
			}

			push(@{$sections->{$level}}, $result->{'secname'});
		}
	}

	die "$0: ERROR: no initcalls?" if (!keys(%{$sections}));

	# print out a linker script that defines the order of initcalls for
	# each level
	print "SECTIONS {\n";

	foreach my $level (sort(keys(%{$sections}))) {
		my $section;

		if ($level eq 'con') {
			$section = '.con_initcall.init';
		} else {
			$section = ".initcall${level}.init";
		}

		print "\t${section} : {\n";

		foreach my $secname (@{$sections->{$level}}) {
			print "\t\t*(${section}..${secname}) ;\n";
		}

		print "\t}\n";
	}

	print "}\n";
}

generate_initcall_lds();