Shorten the generated code with ./release-rs.sh --features short

.cargo/x86_64-unknown-linux-gnu-short.json

@@ -0,0 +1,50 @@
+{
+ "arch": "x86_64",
+ "cpu": "x86-64",
+ "crt-objects-fallback": "false",
+ "crt-static-respected": true,
+ "data-layout": "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128",
+ "dynamic-linking": true,
+ "eh-frame-header": false,
+ "env": "gnu",
+ "has-rpath": true,
+ "has-thread-local": true,
+ "is-builtin": false,
+ "linker-flavor": "gnu-cc",
+ "llvm-target": "x86_64-unknown-linux-gnu",
+ "max-atomic-width": 64,
+ "os": "linux",
+ "plt-by-default": false,
+ "position-independent-executables": true,
+ "pre-link-args": {
+ "gnu-cc": [
+ "-m64"
+ ],
+ "gnu-lld-cc": [
+ "-m64"
+ ]
+ },
+ "relro-level": "off",
+ "stack-probes": {
+ "kind": "none"
+ },
+ "static-position-independent-executables": true,
+ "supported-sanitizers": [
+ "address",
+ "cfi",
+ "leak",
+ "memory",
+ "thread",
+ "safestack"
+ ],
+ "supported-split-debuginfo": [
+ "packed",
+ "unpacked",
+ "off"
+ ],
+ "supports-xray": true,
+ "target-family": [
+ "unix"
+ ],
+ "target-pointer-width": "64"
+}

Cargo.toml

@@ -24,10 +24,21 @@ bench = false
 path = "src/bin/basm-submit.rs"
 
 [dependencies]
-compiler_builtins = { version = "0.1.101", features = ["mem"] }
 libm = "0.2.7"
 ryu = "1.0"
 
+[target.x86_64-pc-windows-msvc.dependencies]
+compiler_builtins = { version = "0.1.103", features = ["mem"] }
+[target.x86_64-unknown-linux-gnu.dependencies]
+compiler_builtins = { version = "0.1.103", features = ["mem"] }
+[target.i686-unknown-linux-gnu.dependencies]
+compiler_builtins = { version = "0.1.103", features = ["mem"] }
+[target.wasm32-unknown-unknown.dependencies]
+compiler_builtins = { version = "0.1.103", features = ["mem"] }
+
+[features]
+short = []
+
 [profile.dev]
 panic = "abort"
 

build.rs

@@ -18,7 +18,7 @@ fn main() {
  link_args_basm.push("/EMITPOGOPHASEINFO");
  link_args_basm_submit.push("/ALIGN:128");
  },
- "x86_64-unknown-linux-gnu" | "i686-unknown-linux-gnu" => {
+ "x86_64-unknown-linux-gnu" | "x86_64-unknown-linux-gnu-short" | "i686-unknown-linux-gnu" => {
  link_args_basm.push("-nostartfiles");
  link_args_basm.push("-nostdlib");
  link_args_basm.push("-static-pie");
@@ -27,12 +27,12 @@ fn main() {
  link_args_basm.push("-fno-unwind-tables");
  link_args_basm.push("-fno-stack-protector");
  link_args_basm.push("-fno-plt");
- if target == "x86_64-unknown-linux-gnu" {
- link_args_basm.push("-Wl,--build-id=none,--gc-sections,--no-eh-frame-hdr,-z,norelro");
- } else {
+ if target == "i686-unknown-linux-gnu" {
  // Prevent linker from putting data into text, which is non-writable and hence not relocatable.
  // This prevents the hack for getting the _DYNAMIC symbol in the entrypoint.
  link_args_basm.push("-Wl,--build-id=none,--gc-sections,--no-eh-frame-hdr,-z,norelro,-z,notext");
+ } else {
+ link_args_basm.push("-Wl,--build-id=none,--gc-sections,--no-eh-frame-hdr,-z,norelro");
  }
  link_args_basm_submit.push("-Wl,-z,max-page-size=128");
  },

scripts/static-pie-prestub-amd64-print.py

@@ -1,5 +1,8 @@
+import sys
+
 # read prestub
-with open("static-pie-prestub-amd64.bin", "rb") as f:
+fname = "static-pie-prestub-amd64.bin" if len(sys.argv) <= 1 else sys.argv[1]
+with open(fname, "rb") as f:
  prestub = f.read()
 prestub = bytearray(prestub)
 if len(prestub) > 0 and prestub[-1] == 0:

scripts/static-pie-prestub-amd64-short.asm

@@ -0,0 +1,86 @@
+; -*- tab-width: 4 -*-
+; 
+; The prestub for amd64-rust target
+; (prestub: the code that runs before the stub and sets the stage)
+;
+; build: nasm -f bin -O9 static-pie-prestub-amd64-short.asm -o static-pie-prestub-amd64-short.bin
+; note: after building with the above command, run static-pie-prestub-amd64-print.py static-pie-prestub-amd64-short.bin
+; to obtain the form that can be embedded in Rust as inline assembly.
+
+BITS 64
+ORG 0
+section .text
+
+; Align stack to 16 byte boundary
+; [rsp+ 32, rsp+120): PLATFORM_DATA
+; [rsp+ 0, rsp+ 32): (shadow space for win64 calling convention)
+ enter 48, 0
+ push 1
+ pop rcx ; Enable ENV_FLAGS_LINUX_STYLE_CHKSTK outside Windows
+ call _t
+
+; svc_alloc_rwx for Windows and Linux
+; rcx = size
+; rdi = pointer to VirtualAlloc (must be supplied before prepending the mov instruction)
+_svc_alloc_rwx:
+ push 9
+ pop rax ; syscall id of x64 mmap
+ jecxz _decode
+ cdq ; rdx=0
+ xor r9d, r9d ; offset
+ push rsi ; save rsi
+ xor edi, edi ; rdi=0
+ mov esi, ecx ; size
+ mov dl, 7 ; protect (safe since we have ensured rdx=0)
+ push 0x22
+ pop r10 ; flags
+ push -1
+ pop r8 ; fd
+ syscall
+ pop rsi ; restore rsi
+_ret:
+ ret
+_svc_alloc_rwx_end:
+
+; Base91 decoder
+_decode:
+ mov al, 0x1f
+_decode_loop:
+ shl eax, 13
+ lodsb
+ sub al, 0x24
+ jc _ret
+ cdq
+ xchg eax, edx
+ lodsb
+ sub al, 0x24
+ imul eax, eax, 91
+ add eax, edx
+_decode_output:
+ stosb
+ shr eax, 8
+ test ah, 16
+ jnz _decode_output
+ jmp _decode_loop
+
+; PLATFORM_DATA
+_t: ; PLATFORM_DATA[32..39] = ptr_alloc_rwx
+ pop rbx
+ push rbx
+ push rdx ; PLATFORM_DATA[24..31] = win_GetProcAddress
+ push rax ; PLATFORM_DATA[16..23] = win_kernel32
+ push rcx ; PLATFORM_DATA[ 8..15] = env_flags (0=None, 1=ENV_FLAGS_LINUX_STYLE_CHKSTK)
+ inc ecx
+ push rcx ; PLATFORM_DATA[ 0.. 7] = env_id (1=Windows, 2=Linux)
+ sub rsp, 40 ; shadow space + compensation
+ call rbx ; svc_alloc_rwx
+
+; Current state: rax = new buffer, rdi = pointer to VirtualAlloc
+ push rax
+ xchg rax, rdi ; rdi = new buffer
+
+; Decode stub (rsi -> rdi)
+; Current state: rdi = stub memory (by the previous instruction)
+; rsi = STUB_BASE91 (by the Rust template)
+ xor ecx, ecx
+ jmp rbx ; This will jump to the start of the new buffer (stub) upon the ret instruction

scripts/static-pie-template-amd64-short.c

@@ -0,0 +1,67 @@
+// Generated with https://github.com/kiwiyou/basm-rs
+// Learn rust (https://doc.rust-lang.org/book/) and get high performance out of the box!
+// SOLUTION BEGIN
+$$$$solution_src$$$$
+// SOLUTION END
+#include <unistd.h>
+typedef unsigned char u8;
+typedef unsigned int u32;
+typedef unsigned long long u64;
+#define BASMCALL __attribute__((ms_abi))
+// Base85 decoder. Code adapted from:
+// https://github.com/rafagafe/base85/blob/master/base85.c
+const char *b85 = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!#$%&()*+-;<=>\?@^_`{|}~";
+void b85tobin(void *dest, char const *src) {
+ u32 *p = (u32 *)dest;
+ u8 digittobin[256];
+ for (u8 i=0; i<85; i++) digittobin[(u8)b85[i]] = i;
+ while (1) {
+ while (*src == '\0') src++;
+ if (*src == ']') break;
+ u32 value = 0;
+ for (u32 i=0; i<5; i++) {
+ value *= 85;
+ value += digittobin[(u8)*src++];
+ }
+ *p++ = (value >> 24) | ((value >> 8) & 0xff00) | ((value << 8) & 0xff0000) | (value << 24);
+ }
+}
+#pragma pack(push, 1)
+typedef struct {
+ u64 env_id;
+ u64 env_flags;
+ u64 win[2];
+ void *fn_table[6];
+} PLATFORM_DATA;
+#pragma pack(pop)
+typedef int (BASMCALL *stub_ptr)(void *, void *);
+char payload[][$$$$min_len_4096$$$$] = $$$$binary_base85$$$$;
+int main() {}
+#ifdef __cplusplus
+extern "C"
+#endif
+int __libc_start_main(
+ void *func_ptr,
+ int argc,
+ char* argv[],
+ void (*init_func)(void),
+ void (*fini_func)(void),
+ void (*rtld_fini_func)(void),
+ void *stack_end) {
+ PLATFORM_DATA pd;
+ pd.env_id = 2;
+ pd.env_flags = 1;
+ u8 stubbuf[68 + $$$$stub_len$$$$];
+ b85tobin(stubbuf, "QMd~L002n8@6D@;XGJ3cz5oya01pLO>naZmS5~+Q0000n|450>x(5IN07=KfA^-pYO)<bp|Hw@-$qxlyU&9Xz]");
+ b85tobin(stubbuf + 68, $$$$stub_base85$$$$);
+ size_t base = ((size_t)main) & 0xFFFFFFFFFFFFF000ULL;
+ *(u64 *)(stubbuf + 0x08) = (u64) base;
+ *(u32 *)(stubbuf + 0x11) = (u32) 4096;
+ base = ((size_t)stubbuf) & 0xFFFFFFFFFFFFF000ULL;
+ size_t len = (((size_t)stubbuf) + 68 + $$$$stub_len$$$$) - base;
+ len = ((len + 0xFFF) >> 12) << 12;
+ syscall(10, base, len, 0x7);
+ pd.fn_table[0] = (void *) (stubbuf + 0x1c);
+ b85tobin(payload, (char const *)payload);
+ return ((stub_ptr) stubbuf)(&pd, payload);
+}

scripts/static-pie-template-amd64-short.rs

@@ -0,0 +1,8 @@
+// Generated with https://github.com/kiwiyou/basm-rs
+// Learn rust (https://doc.rust-lang.org/book/) and get high performance out of the box!
+// SOLUTION BEGIN
+#![crate_type="cdylib"]#![no_std]#[cfg(any())]mod x{
+$$$$solution_src$$$$
+}
+// SOLUTION END
+#[no_link]extern crate std;static mut P:[u8;$$$$binary_base91_len$$$$]=*br$$$$binary_base91$$$$;#[no_mangle]unsafe fn _start(){std::arch::asm!(".quad 0e859016a000030c8h,6758096a0000003ch,3156c931459917e3h,41226a07b2ce89ffh,5e050f5841ff6a5ah,2cac0de0c11fb0c3h,242cac9299f57224h,0e8c1aad0015bc06bh,0e3ebf77510c4f608h,51c1ff515052535bh,4850d3ff28ec8348h,0e3ffc93197h",in("r14")P.as_mut_ptr(),in("rsi")r$$$$stub_base91$$$$.as_ptr())}

scripts/static-pie.sh

@@ -11,9 +11,17 @@ shift
 if [[ "$target_name" == "x86_64-unknown-linux-gnu" ]]; then
  stub="static-pie-stub-amd64.bin"
  if [[ "$lang_name" == "C" ]]; then
- template="static-pie-template-amd64.c"
+ if [[ "$*" == *"short"* ]]; then
+ template="static-pie-template-amd64-short.c"
+ else
+ template="static-pie-template-amd64.c"
+ fi
  elif [[ "$lang_name" == "Rust" ]]; then
- template="static-pie-template-amd64.rs"
+ if [[ "$*" == *"short"* ]]; then
+ template="static-pie-template-amd64-short.rs"
+ else
+ template="static-pie-template-amd64.rs"
+ fi
  else
  >&2 echo "Language ${lang_name} is not supported for target ${target_name}"
  exit
@@ -48,20 +56,30 @@ else
  >&2 echo "Unknown build mode ${build_mode}"
  exit
 fi
+
+if [[ "$target_name" == "x86_64-unknown-linux-gnu" && "$*" == *"short"* ]]; then
+ target_name_cargo=".cargo/x86_64-unknown-linux-gnu-short.json"
+ target_name="x86_64-unknown-linux-gnu-short"
+ extra_config='-Zbuild-std=core,compiler_builtins,alloc -Zbuild-std-features=compiler-builtins-mem'
+else
+ target_name_cargo="$target_name"
+ extra_config=""
+fi
+
 >&2 echo "Building project for target ${target_name}, language ${lang_name}, build mode ${build_mode}"
 
 binary_path=basm.bin
 if [[ "$build_mode" == "Debug" ]]; then
- cargo +nightly build --target "$target_name" --bin basm-submit "$@"
+ cargo +nightly build $extra_config --target "$target_name_cargo" --bin basm-submit "$@"
 else
- cargo +nightly build --target "$target_name" --bin basm-submit --release "$@"
+ cargo +nightly build $extra_config --target "$target_name_cargo" --bin basm-submit --release "$@"
 fi
 
 if [[ "$target_name" == "x86_64-pc-windows-msvc" ]]; then
  python3 scripts/static-pie-gen.py src/solution.rs "$target_name" target/"$target_name"/"$build_mode_dir"/basm-submit.exe scripts/"$stub" "$lang_name" scripts/"$template"
 else
  cp target/"$target_name"/"$build_mode_dir"/basm-submit target/"$target_name"/"$build_mode_dir"/basm-submit-stripped
  objcopy --strip-all target/"$target_name"/"$build_mode_dir"/basm-submit-stripped
- objcopy --remove-section .eh_frame target/"$target_name"/"$build_mode_dir"/basm-submit-stripped
+ objcopy --remove-section .eh_frame --remove-section .gcc_except_table --remove-section .gnu.hash target/"$target_name"/"$build_mode_dir"/basm-submit-stripped
  python3 scripts/static-pie-gen.py src/solution.rs "$target_name" target/"$target_name"/"$build_mode_dir"/basm-submit-stripped scripts/"$stub" "$lang_name" scripts/"$template"
 fi

src/bin/codegen.rs

@@ -59,6 +59,7 @@ unsafe extern "win64" fn _start() -> ! {
  // on the 16-byte boundary BEFORE `call` instruction.
  // However, when called as the entrypoint by the Linux OS,
  // RSP will be 16-byte aligned AFTER `call` instruction.
+ #[cfg(not(feature = "short"))]
  asm!(
  "clc", // CF=0 (running without loader) / CF=1 (running with loader)
  "mov rbx, rcx", // Save PLATFORM_DATA table
@@ -80,6 +81,22 @@ unsafe extern "win64" fn _start() -> ! {
  sym _start_rust,
  options(noreturn)
  );
+ // For "short", we always assume we are running with loader on Linux,
+ // since "short" is only meaningful when submitting to online judges (not local test runs).
+ #[cfg(feature = "short")]
+ asm!(
+ "clc", // Not needed but packager wants it
+ "push rcx", // Align stack
+ "mov rbx, rcx", // Save PLATFORM_DATA table
+ "lea rdi, [rip + __ehdr_start]",
+ "lea rsi, [rip + _DYNAMIC]",
+ "call {0}",
+ "mov rdi, rbx",
+ "call {1}", // This won't return since on Linux we invoke SYS_exitgroup in binary
+ sym loader::amd64_elf::relocate,
+ sym _start_rust,
+ options(noreturn)
+ );
 }
 
 #[cfg(target_os = "windows")]
@@ -234,11 +251,11 @@ extern "C" fn _start() {
 }
 
 /* We prevent inlining solution::main, since if the user allocates
- * a large amount of stack memory there, it will be zero-initialized
+ * a large amount of stack memory there, it will be zero-initialized (or probed)
  * *before* we increase the stack limits if it is inlined into _start_rust.
  * This will cause stack overflow, thus we prevent it.
  */
-#[inline(never)]
+#[cfg_attr(not(feature = "short"), inline(never))]
 fn _call_main() {
  solution::main();
 }

src/lib.rs

@@ -1,3 +1,4 @@
+#![feature(rustc_private)]
 #![feature(fn_align)]
 #![feature(maybe_uninit_slice)]
 #![feature(maybe_uninit_uninit_array)]