######################## capture.cpp

#pragma once

#include <string>
#include <string_view>
#include <utility>
#include <variant>

#include <unused.hpp>  // TODO: remove before flight.

template <class F>
std::variant<std::pair<std::string, std::string>, int>
CaptureOutput(F&& f, std::string_view input) {
    UNUSED(f, input);  // TODO: remove before flight.
    // TODO: your code here.
    throw "TODO";
}


######################## distribution.hpp

#pragma once

#include <random>

struct UniformCharDistribution {
    using result_type = char;

    UniformCharDistribution(char lhs, char rhs)
        : impl_(static_cast<int>(lhs), static_cast<int>(rhs)) {
    }

    template <class Rng>
    char operator()(Rng& rng) {
        return static_cast<char>(impl_(rng));
    }

  private:
    std::uniform_int_distribution<int> impl_;
};


######################## fd-guard.cpp

#include "fd-guard.hpp"
#include <cassert>
#include <algorithm>
#include <filesystem>
#include <unistd.h>

#include <fcntl.h>
#include <linux/kcmp.h>
#include <sys/resource.h>
#include <sys/syscall.h>

#define WARN(v)

namespace fs = std::filesystem;

static std::vector<int> GetOpenFileDescriptors() {
    std::vector<int> fds_candidates;
    for (auto dent : fs::directory_iterator{"/proc/self/fd"}) {
        fds_candidates.push_back(std::stoi(dent.path().filename()));
    }
    std::vector<int> fds;
    for (auto fd : fds_candidates) {
        int ret = fcntl(fd, F_GETFD);
        if (ret == -1) {
            continue;
        }
        fds.push_back(fd);
    }
    std::sort(fds.begin(), fds.end());

    return fds;
}

static size_t MaxFdNum() {
    struct rlimit limits{};

    int r = ::getrlimit(RLIMIT_NOFILE, &limits);
    assert(r != -1);
    // Probably should check `/proc/sys/fs/nr_open` instead
    return std::min(static_cast<size_t>(limits.rlim_cur), size_t{2048});
}

FileDescriptorsGuard::FileDescriptorsGuard() {
    auto fds = GetOpenFileDescriptors();
    if (fds.empty()) {
        return;
    }

    auto max_fd_num = MaxFdNum();
    assert(2 * fds.size() < max_fd_num);

    auto cur_fd_num = static_cast<int>(max_fd_num - 1);
    for (auto fd : fds) {
        while (std::ranges::find(fds, cur_fd_num) != fds.end()) {
            --cur_fd_num;
        }

        int ret = dup3(fd, cur_fd_num, O_CLOEXEC);
        assert(ret != -1);
        copies.emplace_back(fd, cur_fd_num);
        --cur_fd_num;
    }
}

bool FileDescriptorsGuard::TestDescriptorsState() const {
    for (auto [from, to] : copies) {
        if (!TestSameFd(from, to)) {
            WARN("Descriptors " << from << " and " << to
                                << " don't refer to the same file");
            return false;
        }
    }

    auto fds = GetOpenFileDescriptors();
    auto expected = OpenFdsCount();
    if (fds.size() != expected) {
        WARN("Descriptors count mismatch. Expected "
             << expected << ", actually " << fds.size());
        return false;
    }
    return true;
}

size_t FileDescriptorsGuard::OpenFdsCount() const {
    return copies.size() * 2;
}

bool FileDescriptorsGuard::TestSameFd(int fd1, int fd2) const {
    int pid = getpid();
    int ret = syscall(SYS_kcmp, pid, pid, KCMP_FILE, fd1, fd2);
    if (ret == -1) {
        int err = errno;
        if (err == ENOSYS) {
            if (ShouldWarnKCmp()) {
                WARN("KCmp is not available on your system, some checks "
                     "wouldn't be performed");
            }
            return true;
        }
        assert(err == EBADF);
        return false;
    }
    return ret == 0;
}

bool FileDescriptorsGuard::ShouldWarnKCmp() const {
    return !std::exchange(warned_kcmp_, true);
}

FileDescriptorsGuard::~FileDescriptorsGuard() {
    for (auto [_, fd] : copies) {
        close(fd);
    }
}


######################## fd-guard.hpp

#pragma once

#ifndef __linux__
#error "Only linux is supported for this problem"
#endif

#include <cstddef>
#include <utility>
#include <vector>

struct FileDescriptorsGuard {
    FileDescriptorsGuard();

    [[nodiscard]] bool TestDescriptorsState() const;

    size_t OpenFdsCount() const;

    ~FileDescriptorsGuard();

  private:
    bool ShouldWarnKCmp() const;
    bool TestSameFd(int fd1, int fd2) const;

    std::vector<std::pair<int, int>> copies;
    mutable bool warned_kcmp_{false};
};

######################## overload.hpp

#pragma once

template <class... Ts>
struct Overload : Ts... {
    using Ts::operator()...;
};

template <class... Ts>
Overload(Ts...) -> Overload<Ts...>;


######################## rim-guard.cpp

#include <cassert>
#include "rlim-guard.hpp"

#include <sys/resource.h>
#include <utility>

static uint64_t ModifySoftLimit(int limit_type, uint64_t new_limit) {
    struct rlimit limit{};
    {
        int ret = getrlimit(limit_type, &limit);
        assert(ret != -1);
    }

    auto new_limit_inner = static_cast<rlim_t>(new_limit);
    assert(new_limit_inner <= limit.rlim_max);
    auto prev_limit = std::exchange(limit.rlim_cur, new_limit_inner);

    {
        int ret = setrlimit(limit_type, &limit);
        assert(ret != -1);
    }
    return static_cast<uint64_t>(prev_limit);
}

RLimGuard::RLimGuard(int limit_type, rlim_t new_limit)
    : limit_type_(limit_type) {
    prev_limit_ = ModifySoftLimit(limit_type_, new_limit);
}

void RLimGuard::Reset() {
    auto limit_type = std::exchange(limit_type_, -1);
    if (limit_type == -1) {
        return;
    }
    ModifySoftLimit(limit_type, prev_limit_);
}

RLimGuard::~RLimGuard() {
    Reset();
}

######################## rim-guard.hpp


#pragma once

#include <cstdint>

struct RLimGuard {
    RLimGuard(int limit_type, uint64_t new_limit);

    RLimGuard(const RLimGuard&) = delete;
    RLimGuard(RLimGuard&&) = delete;

    RLimGuard& operator=(const RLimGuard&) = delete;
    RLimGuard& operator=(RLimGuard&&) = delete;

    void Reset();

    ~RLimGuard();

  private:
    int limit_type_;
    uint64_t prev_limit_;
};

######################## test.cpp

#include "capture.hpp"

#include <cassert>
#include "distributions.hpp"
#include "fd-guard.hpp"
#include "overload.hpp"
#include "rlim-guard.hpp"

#include <algorithm>
#include <iostream>
#include <random>

#include <fcntl.h>
#include <sys/resource.h>
#include <unistd.h>

#define TEST_CASE(v)
#define SECTION(v)
#define INFO(v)
#define REQUIRE(v) assert(v)
#define CHECK(v) assert(v)


void Flush() {
    std::cout.flush();
    std::clog.flush();
    std::fflush(stdout);
    std::fflush(stderr);
}

template <class Rng>
std::string GenerateStr(Rng& rng, size_t n) {
    UniformCharDistribution dist('a', 'z');
    std::string s(n, ' ');
    std::generate(s.begin(), s.end(), [&rng, &dist] { return dist(rng); });
    return s;
}

void ResetCinState() {
    std::cin.clear();
    std::clearerr(stdin);
}

static constexpr size_t kPipeSize = 4 << 10;

int main() {
TEST_CASE("CheckPipeCapacity") {
    INFO("This is an internal assertion, please report if it fails");
    int fds[2];
    int r = pipe(fds);
    REQUIRE(r != -1);
    int cap = fcntl(fds[0], F_GETPIPE_SZ);
    {
        int err = errno;
        INFO("Errno is " << err);
        REQUIRE(cap != -1);
    }
    REQUIRE(static_cast<size_t>(cap) >= kPipeSize);
    close(fds[0]);
    close(fds[1]);
}

TEST_CASE("JustWorks") {
    FileDescriptorsGuard guard;

    SECTION("Simple") {
        int runs = 0;
        Flush();
        auto result = CaptureOutput([&runs] { ++runs; }, "");
        REQUIRE(result.index() == 0);
        auto [out, err] = std::get<0>(std::move(result));
        CHECK(out == "");
        CHECK(err == "");
        CHECK(runs == 1);
        CHECK(guard.TestDescriptorsState());
    }

    SECTION("Output") {
        Flush();
        auto result = CaptureOutput(
            [] {
                std::cout << "Aba";
                std::cout.flush();
                std::cerr << "Caba";
            },
            "");
        REQUIRE(result.index() == 0);
        auto [out, err] = std::get<0>(std::move(result));
        CHECK(out == "Aba");
        CHECK(err == "Caba");
        CHECK(guard.TestDescriptorsState());
    }

    SECTION("Input+Output") {
        Flush();
        auto result = CaptureOutput(
            [] {
                for (size_t i = 0; i < 3; ++i) {
                    int v;
                    std::cin >> v;
                    if (v & 1) {
                        std::cout << v << " ";
                    } else {
                        std::cerr << v << " ";
                    }
                }

                std::cout.flush();
            },
            "1 2 3 ");

        REQUIRE(result.index() == 0);
        auto [out, err] = std::get<0>(std::move(result));
        CHECK(out == "1 3 ");
        CHECK(err == "2 ");
        CHECK(guard.TestDescriptorsState());
    }
}

TEST_CASE("EOF") {
    FileDescriptorsGuard guard;

    Flush();
    auto result = CaptureOutput(
        [] {
            int v;
            while (std::cin >> v) {
                if (v & 1) {
                    std::cout << v;
                } else {
                    std::cerr << v;
                }
            }
            std::cout.flush();
        },
        "1 2  3 4 5  6 7 8");
    ResetCinState();

    REQUIRE(result.index() == 0);
    auto [out, err] = std::get<0>(std::move(result));
    CHECK(out == "1357");
    CHECK(err == "2468");
    CHECK(guard.TestDescriptorsState());
}

TEST_CASE("HugeIO") {
    FileDescriptorsGuard guard;

    static constexpr size_t kBufSize = kPipeSize;
    std::mt19937 rng(42);

    auto inp = GenerateStr(rng, kBufSize);
    auto my_out = GenerateStr(rng, kBufSize);
    auto my_err = GenerateStr(rng, kBufSize);

    std::string actual_inp;
    Flush();
    auto result = CaptureOutput(
        [&] {
            std::string s;
            std::cin >> s;

            std::cout << my_out;
            std::cerr << my_err;

            actual_inp = std::move(s);
            std::cout.flush();
        },
        inp);
    ResetCinState();

    CHECK(actual_inp == inp);

    REQUIRE(result.index() == 0);
    auto [out, err] = std::get<0>(std::move(result));
    CHECK(out == my_out);
    CHECK(err == my_err);
    CHECK(guard.TestDescriptorsState());
}

TEST_CASE("ErrorRecovery") {
    std::mt19937 rng(42);

    FileDescriptorsGuard guard;
    constexpr auto base_fd_count = 3;

    for (size_t i = 0; i <= 10; ++i) {
        INFO("i = " << i);

        auto out = GenerateStr(rng, 10);
        auto err = GenerateStr(rng, 10);
        auto inp = GenerateStr(rng, 10);

        Flush();
        auto result = [&] {
            RLimGuard files_guard(RLIMIT_NOFILE, base_fd_count + i);
            return CaptureOutput(
                [&] {
                    (std::cout << out).flush();
                    std::cerr << err;
                },
                inp);
        }();

        if (i < 2) {
            REQUIRE(result.index() == 1);
        }
        if (i > 9) {
            REQUIRE(result.index() == 0);
        }

        std::visit(Overload{
                       [&](std::pair<std::string, std::string> output) {
                           CHECK(output.first == out);
                           CHECK(output.second == err);
                       },
                       [](int err) {
                           INFO("Error code: " << err << " "
                                               << std::strerror(err));
                           REQUIRE((err == EBADF || err == EMFILE));
                       },
                   },
                   std::move(result));

        CHECK(guard.TestDescriptorsState());
    }
}

TEST_CASE("RawIO") {
    constexpr size_t kBufSize = 10;

    std::mt19937 rng(42);
    auto input = GenerateStr(rng, kBufSize);
    auto output = GenerateStr(rng, kBufSize);
    auto error = GenerateStr(rng, kBufSize);

    auto write_all = [](int fd, std::string_view data) -> int {
        while (!data.empty()) {
            int w = write(fd, data.data(), data.size());
            if (w == -1) {
                return -errno;
            }
            data = data.substr(w);
        }
        return 0;
    };

    char real_input[kBufSize + 1];

    int out_err = 0;
    int err_err = 0;
    int in_err = 0;
    auto result = CaptureOutput(
        [&] {
            out_err = write_all(STDOUT_FILENO, output);
            err_err = write_all(STDERR_FILENO, error);

            char* in = real_input;
            char* real_in_end = real_input + kBufSize;
            while (in < real_in_end) {
                int r = read(0, in, real_in_end - in);
                if (r == -1) {
                    in_err = -errno;
                    break;
                }
                if (r == 0) {
                    in_err = 1023;
                    break;
                }
                in += r;
            }

            if (read(0, in, 1) != 0) {
                in_err = 1024;
            }
        },
        input);

    REQUIRE(out_err == 0);
    REQUIRE(err_err == 0);
    REQUIRE(in_err == 0);

    REQUIRE(result.index() == 0);
    auto [out, err] = std::get<0>(result);

    CHECK(out == output);
    CHECK(err == error);
    auto val = std::string_view{real_input, kBufSize};
    assert(val == input);
}
}

######################## make.sh

#!/bin/bash
SCR=solution
OUT=runnable

cp $filename capture.hpp

g++ fd-guard.cpp rim-guard.cpp test.cpp -fsanitize=address -std=c++20 -o $OUT