/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use crate::allocative_trait::Allocative;
use crate::visitor::Visitor;

impl Allocative for u8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for u16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for u32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for u64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for usize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for i8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for i16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for i32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for i64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for isize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for bool {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for f32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for f64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::mem;
use std::rc;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI16;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::AtomicI8;
use std::sync::atomic::AtomicIsize;
use std::sync::atomic::AtomicU16;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use std::sync::Weak;

use crate::allocative_trait::Allocative;
use crate::impls::common::PTR_NAME;
use crate::key::Key;
use crate::visitor::Visitor;

impl<T: Allocative> Allocative for RwLock<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if let Ok(data) = self.try_read() {
            visitor.visit_field(Key::new("data"), &*data);
        }
        visitor.exit();
    }
}

impl<T: Allocative + ?Sized> Allocative for Arc<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let visitor = visitor.enter_shared(
                PTR_NAME,
                // TODO(nga): 8 or 16 depending on sized or unsized.
                mem::size_of::<*const ()>(),
                Arc::as_ptr(self) as *const (),
            );
            if let Some(mut visitor) = visitor {
                struct ArcInner(AtomicUsize, AtomicUsize, ());
                {
                    let val: &T = self;
                    let mut visitor = visitor.enter(
                        Key::new("ArcInner"),
                        mem::size_of::<ArcInner>() + mem::size_of_val(val),
                    );
                    val.visit(&mut visitor);
                    visitor.exit();
                }
                visitor.exit();
            }
        }
        visitor.exit();
    }
}

impl<T: Allocative + ?Sized> Allocative for Weak<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            if let Some(arc) = self.upgrade() {
                arc.visit(&mut visitor);
            }
        }
        visitor.exit();
    }
}

impl<T: Allocative> Allocative for Rc<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let visitor = visitor.enter_shared(
                PTR_NAME,
                // TODO(nga): 8 or 16 depending on sized or unsized.
                mem::size_of::<*const ()>(),
                Rc::as_ptr(self) as *const (),
            );
            if let Some(mut visitor) = visitor {
                struct RcInner(AtomicUsize, AtomicUsize, ());
                {
                    let val: &T = self;
                    let mut visitor = visitor.enter(
                        Key::new("RcInner"),
                        mem::size_of::<RcInner>() + mem::size_of_val(val),
                    );
                    val.visit(&mut visitor);
                    visitor.exit();
                }
                visitor.exit();
            }
        }
        visitor.exit();
    }
}

impl<T: Allocative> Allocative for rc::Weak<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            if let Some(rc) = self.upgrade() {
                rc.visit(&mut visitor);
            }
        }
        visitor.exit();
    }
}

impl Allocative for AtomicU8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicU16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicU32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicU64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicUsize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicI8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicI16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicI32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicI64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicBool {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for AtomicIsize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl<T: Allocative> Allocative for Mutex<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if let Ok(data) = self.try_lock() {
            visitor.visit_field(Key::new("data"), &*data);
        }
        visitor.exit();
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;

use crate::allocative_trait::Allocative;
use crate::visitor::Visitor;

impl Allocative for Instant {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for SystemTime {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}

impl Allocative for Duration {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.enter_self_sized::<Self>().exit();
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use crate::allocative_trait::Allocative;
use crate::key::Key;
use crate::visitor::Visitor;

impl Allocative for () {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl<A: Allocative> Allocative for (A,) {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(Key::new("0"), &self.0);
    }
}

impl<A: Allocative, B: Allocative> Allocative for (A, B) {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(Key::new("0"), &self.0);
        visitor.visit_field(Key::new("1"), &self.1);
    }
}

impl<A: Allocative, B: Allocative, C: Allocative> Allocative for (A, B, C) {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(Key::new("0"), &self.0);
        visitor.visit_field(Key::new("1"), &self.1);
        visitor.visit_field(Key::new("2"), &self.2);
    }
}

impl<A: Allocative, B: Allocative, C: Allocative, D: Allocative> Allocative for (A, B, C, D) {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(Key::new("0"), &self.0);
        visitor.visit_field(Key::new("1"), &self.1);
        visitor.visit_field(Key::new("2"), &self.2);
        visitor.visit_field(Key::new("3"), &self.3);
    }
}

impl<A: Allocative, B: Allocative, C: Allocative, D: Allocative, E: Allocative> Allocative
    for (A, B, C, D, E)
{
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        visitor.visit_field(Key::new("0"), &self.0);
        visitor.visit_field(Key::new("1"), &self.1);
        visitor.visit_field(Key::new("2"), &self.2);
        visitor.visit_field(Key::new("3"), &self.3);
        visitor.visit_field(Key::new("4"), &self.4);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::VecDeque;
use std::convert::Infallible;
use std::ffi::OsStr;
use std::marker::PhantomData;
use std::mem;
use std::num::NonZeroI16;
use std::num::NonZeroI32;
use std::num::NonZeroI64;
use std::num::NonZeroI8;
use std::num::NonZeroIsize;
use std::num::NonZeroU16;
use std::num::NonZeroU32;
use std::num::NonZeroU64;
use std::num::NonZeroU8;
use std::num::NonZeroUsize;
use std::path::PathBuf;

use crate::allocative_trait::Allocative;
use crate::impls::common::DATA_NAME;
use crate::impls::common::PTR_NAME;
use crate::impls::common::UNUSED_CAPACITY_NAME;
use crate::key::Key;
use crate::visitor::Visitor;

impl<T: Allocative + ?Sized> Allocative for &'static T {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let _ = visitor;
    }
}

impl Allocative for str {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple(Key::new("str"), self.len());
    }
}

impl Allocative for OsStr {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple(Key::new("OsStr"), self.len());
    }
}

impl<T: Allocative + ?Sized> Allocative for Box<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor = visitor.enter_unique(PTR_NAME, mem::size_of::<*const T>());
            (**self).visit(&mut visitor);
        }
        visitor.exit();
    }
}

impl<T: Allocative, E: Allocative> Allocative for Result<T, E> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        match self {
            Ok(v) => visitor.visit_field(Key::new("Ok"), v),
            Err(e) => visitor.visit_field(Key::new("Err"), e),
        }
    }
}

impl<T: Allocative> Allocative for Option<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if let Some(value) = self {
            visitor.visit_field(Key::new("Some"), value);
        }
    }
}

impl<T: Allocative> Allocative for Vec<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if self.capacity() != 0 && mem::size_of::<T>() != 0 {
            let mut visitor = visitor.enter_unique(PTR_NAME, mem::size_of::<*const T>());
            visitor.visit_slice(self.as_slice());
            visitor.visit_simple(
                UNUSED_CAPACITY_NAME,
                (self.capacity() - self.len()) * mem::size_of::<T>(),
            );
        }
        visitor.exit();
    }
}

impl<T: Allocative> Allocative for VecDeque<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if self.capacity() != 0 && mem::size_of::<T>() != 0 {
            let mut visitor = visitor.enter_unique(PTR_NAME, mem::size_of::<*const T>());
            visitor.visit_slice(self.as_slices().0);
            visitor.visit_slice(self.as_slices().1);
            visitor.visit_simple(
                UNUSED_CAPACITY_NAME,
                (self.capacity() - self.len()) * mem::size_of::<T>(),
            );
        }
        visitor.exit();
    }
}

impl Allocative for String {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        {
            let mut visitor = visitor.enter_unique(PTR_NAME, mem::size_of::<*const u8>());
            visitor.visit_vec_like_body(self.as_bytes(), self.capacity());
            visitor.exit();
        }
        visitor.exit();
    }
}

impl Allocative for PathBuf {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter_self_sized::<Self>();
        if self.capacity() != 0 {
            let mut visitor = visitor.enter_unique(PTR_NAME, mem::size_of::<*const u8>());
            visitor.visit_simple(Key::new("path"), self.as_os_str().len());
            visitor.visit_simple(
                UNUSED_CAPACITY_NAME,
                self.capacity() - self.as_os_str().len(),
            );
        }
        visitor.exit();
    }
}

impl Allocative for NonZeroU64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroU32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroU16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroU8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroI64 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroI32 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroI16 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroI8 {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroUsize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl Allocative for NonZeroIsize {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        visitor.visit_simple_sized::<Self>();
    }
}

impl<T: ?Sized> Allocative for PhantomData<T> {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let _ = visitor;
    }
}

impl Allocative for Infallible {
    fn visit<'a, 'b: 'a>(&self, _visitor: &'a mut Visitor<'b>) {
        match *self {}
    }
}

#[cfg(rust_nightly)]
impl Allocative for ! {
    fn visit<'a, 'b: 'a>(&self, _visitor: &'a mut Visitor<'b>) {
        match *self {}
    }
}

impl<T: Allocative> Allocative for [T] {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        let mut visitor = visitor.enter(Key::for_type_name::<T>(), mem::size_of_val::<[T]>(self));
        for item in self {
            // TODO(nga): faster for primitive.
            visitor.visit_field(DATA_NAME, item);
        }
        visitor.exit();
    }
}

impl<T: Allocative, const N: usize> Allocative for [T; N] {
    fn visit<'a, 'b: 'a>(&self, visitor: &'a mut Visitor<'b>) {
        self.as_slice().visit(visitor);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use allocative_derive::Allocative;

use crate as allocative;

#[derive(Allocative)]
struct Dst(str);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![cfg(test)]
#![allow(dead_code)]

mod dst;
mod skip;
mod visit;
mod with_flamegraph;

use crate as allocative;
use crate::Allocative;

#[derive(Allocative)]
struct Empty {}

#[derive(Allocative)]
struct TupleStruct(u32, String);

#[derive(Allocative)]
struct RegularStruct {
    a: u32,
    b: String,
}

#[derive(Allocative)]
enum Enum {
    Unit,
    Tuple(u32, String),
    Regular { a: u32, b: String },
}

#[derive(Allocative)]
enum GenericEnum<T> {
    Unit,
    Tuple(T, String),
}

#[derive(Allocative)]
struct StructWithDefaultParam<T = String>(T);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use crate as allocative;
use crate::Allocative;

struct Unsupported;

#[derive(Allocative)]
struct TestIgnoreField {
    #[allocative(skip)]
    _unsupported: Unsupported,
}

#[derive(Allocative)]
#[allocative(skip)]
struct TestSkipOnStruct {
    _unsupported: Unsupported,
}

#[derive(Allocative)]
#[allocative(skip)]
enum TestSkipOnEnum {
    Unsupported(Unsupported),
}

#[derive(Allocative)]
struct TestIgnoreInTupleStruct(#[allocative(skip)] Unsupported);

#[derive(Allocative)]
enum TestAllocativeIgnoreEnumVariant {
    String(String),
    #[allocative(skip)]
    Unsupported(Unsupported),
    UnsupportedInside(#[allocative(skip)] Unsupported),
}

#[derive(Allocative)]
#[allocative(bound = "")]
struct IgnoreBound<T> {
    #[allocative(skip)]
    t: T,
}

#[derive(Allocative)]
struct TestBoundIgnored {
    ignore_bound: IgnoreBound<Unsupported>,
}

#[derive(Allocative)]
#[allocative(skip)]
struct TypeLevelSkipImpliesNoBounds<T> {
    t: T,
}

#[derive(Allocative)]
struct TestTypeLevelSkipImpliesNoBounds {
    type_level_skip: TypeLevelSkipImpliesNoBounds<Unsupported>,
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use allocative_derive::Allocative;

use crate as allocative;
use crate::Visitor;

struct DoesNotImplementAllocative<T>(T);

fn visit_does_not_implement_allocative<T>(
    _field: &DoesNotImplementAllocative<T>,
    visitor: &mut Visitor<'_>,
) {
    visitor
        .enter_self_sized::<DoesNotImplementAllocative<T>>()
        .exit();
}

#[derive(Allocative)]
struct MyStruct {
    #[allocative(visit = visit_does_not_implement_allocative)]
    field: DoesNotImplementAllocative<String>,
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::mem;

use allocative::Allocative;

use crate as allocative;
use crate::FlameGraphBuilder;

#[derive(Allocative)]
struct TestData {
    data: Box<[u8]>,
    b: u8,
}

#[test]
fn test_flamegraph() {
    if mem::size_of::<usize>() != 8 {
        return;
    }

    let mut fg = FlameGraphBuilder::default();
    fg.visit_root(&TestData {
        data: vec![0; 100].into_boxed_slice(),
        b: 1,
    });
    assert_eq!(
        "\
        allocative::test_derive::with_flamegraph::TestData 7\n\
        allocative::test_derive::with_flamegraph::TestData;b;u8 1\n\
        allocative::test_derive::with_flamegraph::TestData;data;alloc::boxed::Box<[u8]>;ptr 16\n\
        allocative::test_derive::with_flamegraph::TestData;data;alloc::boxed::Box<[u8]>;ptr;u8;data;u8 100\n\
        ",
        // When running test with buck, crate name is `allocative_unittest`.
        fg.finish_and_write_flame_graph()
            .replace("allocative_unittest::", "allocative::")
    );
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use proc_macro::TokenStream;
use proc_macro2::Ident;
use quote::quote_spanned;
use quote::ToTokens;
use syn::parse::ParseStream;
use syn::spanned::Spanned;
use syn::Attribute;
use syn::Data;
use syn::DataEnum;
use syn::DataStruct;
use syn::DeriveInput;
use syn::Field;
use syn::Fields;
use syn::FieldsNamed;
use syn::FieldsUnnamed;
use syn::GenericParam;
use syn::Generics;
use syn::LitStr;
use syn::Path;
use syn::Token;
use syn::Variant;

const fn hash(s: &str) -> u64 {
    let mut hash = 0xcbf29ce484222325;
    let mut i = 0;
    while i < s.as_bytes().len() {
        let b = s.as_bytes()[i];
        hash ^= b as u64;
        hash = hash.wrapping_mul(0x100000001b3);
        i += 1;
    }
    hash
}

pub(crate) fn derive_allocative(input: TokenStream) -> TokenStream {
    match derive_allocative_impl(input.into()) {
        Ok(tokens) => tokens.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

fn impl_generics(
    generics: &Generics,
    attrs: &AllocativeAttrs,
) -> syn::Result<proc_macro2::TokenStream> {
    if let Some(bound) = &attrs.bound {
        if !bound.is_empty() {
            return Err(syn::Error::new(
                attrs.bound.span(),
                "non-empty bound is not implemented",
            ));
        }
    }

    let mut impl_generics = Vec::new();
    for p in &generics.params {
        impl_generics.push(match p {
            GenericParam::Type(tp) => {
                let mut tp = tp.clone();
                if attrs.bound.is_none() && !attrs.skip {
                    tp.bounds.push(syn::parse2(quote_spanned! { tp.span() =>
                        allocative::Allocative
                    })?);
                }
                tp.default = None;
                tp.to_token_stream()
            }
            GenericParam::Lifetime(l) => l.to_token_stream(),
            GenericParam::Const(c) => c.to_token_stream(),
        });
    }
    if impl_generics.is_empty() {
        Ok(quote_spanned! { generics.span() => })
    } else {
        Ok(quote_spanned! { generics.span() => <#(#impl_generics),*> })
    }
}

fn derive_allocative_impl(
    input: proc_macro2::TokenStream,
) -> syn::Result<proc_macro2::TokenStream> {
    let input: DeriveInput = syn::parse2(input)?;
    let name = &input.ident;
    let (_impl_generics, type_generics, where_clause) = input.generics.split_for_impl();

    let attrs = extract_attrs(&input.attrs)?;
    let impl_generics = impl_generics(&input.generics, &attrs)?;

    let body = if attrs.skip {
        quote_spanned! { input.span() =>
        }
    } else {
        gen_visit_body(&input)?
    };

    Ok(quote_spanned! {input.span()=>
        impl #impl_generics allocative::Allocative for #name #type_generics #where_clause {
            #[allow(unused, warnings)]
            fn visit<'allocative_a, 'allocative_b: 'allocative_a>(
                &self,
                visitor: &'allocative_a mut allocative::Visitor<'allocative_b>,
            ) {
                let mut visitor = visitor.enter_self::<Self>(self);
                #body
                visitor.exit();
            }
        }
    })
}

fn gen_visit_body(input: &DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
    match &input.data {
        Data::Struct(s) => gen_visit_struct(s),
        Data::Union(_) => Err(syn::Error::new_spanned(input, "unions are not supported")),
        Data::Enum(e) => gen_visit_enum(e),
    }
}

fn gen_visit_enum(input: &DataEnum) -> syn::Result<proc_macro2::TokenStream> {
    let cases = input
        .variants
        .iter()
        .map(gen_visit_enum_variant)
        .collect::<syn::Result<proc_macro2::TokenStream>>()?;
    Ok(quote_spanned! {input.variants.span()=>
        match self {
            #cases
        }
    })
}

fn allocative_key(s: &str) -> proc_macro2::TokenStream {
    // Compile hash at proc macro time, otherwise it will have to be computed by MIRI.
    let hash = hash(s);
    quote_spanned! {proc_macro2::Span::call_site()=>
        allocative::Key::new_unchecked(#hash, #s)
    }
}

fn gen_visit_enum_variant(input: &Variant) -> syn::Result<proc_macro2::TokenStream> {
    let name = &input.ident;
    let name_str = name.to_string();

    let variant_attrs = extract_attrs(&input.attrs)?;

    let variant_key = allocative_key(&name_str);

    // TODO: enter variant.
    match &input.fields {
        Fields::Unit => Ok(quote_spanned! {input.span()=>
            Self::#name => {},
        }),
        Fields::Unnamed(unnamed) => {
            let field_names = fields_unnamed_names(unnamed)?;
            let visit_fields = unnamed
                .unnamed
                .iter()
                .enumerate()
                .zip(field_names.iter())
                .map(|((i, f), n)| {
                    if variant_attrs.skip {
                        Ok(quote_spanned! {f.span()=>
                        })
                    } else {
                        gen_visit_field(&i.to_string(), n, f)
                    }
                })
                .collect::<syn::Result<proc_macro2::TokenStream>>()?;
            Ok(quote_spanned! {input.span()=>
                Self::#name(#(#field_names),*) => {
                    let mut visitor = visitor.enter(#variant_key, std::mem::size_of::<Self>());
                    #visit_fields
                    visitor.exit();
                },
            })
        }
        Fields::Named(named) => {
            let field_names = fields_named_names(named)?;
            let visit_fields = field_names
                .iter()
                .zip(named.named.iter())
                .map(|(ident, f)| {
                    if variant_attrs.skip {
                        Ok(quote_spanned! {f.span()=>
                        })
                    } else {
                        gen_visit_field(&ident.to_string(), ident, f)
                    }
                })
                .collect::<syn::Result<proc_macro2::TokenStream>>()?;
            Ok(quote_spanned! {input.span()=>
                Self::#name { #(#field_names),* } => {
                    let mut visitor = visitor.enter(#variant_key, std::mem::size_of::<Self>());
                    #visit_fields
                    visitor.exit();
                },
            })
        }
    }
}

fn fields_named_names(fields: &FieldsNamed) -> syn::Result<Vec<&Ident>> {
    fields
        .named
        .iter()
        .map(|f| {
            f.ident
                .as_ref()
                .ok_or_else(|| syn::Error::new_spanned(f, "unnamed field in named"))
        })
        .collect()
}

fn fields_unnamed_names(fields: &FieldsUnnamed) -> syn::Result<Vec<Ident>> {
    fields
        .unnamed
        .iter()
        .enumerate()
        .map(|(i, f)| Ok(Ident::new(&format!("f{}", i), f.span())))
        .collect()
}

fn gen_visit_struct(input: &DataStruct) -> syn::Result<proc_macro2::TokenStream> {
    match &input.fields {
        Fields::Named(named) => {
            let names = fields_named_names(named)?;
            let visit_fields = names
                .iter()
                .zip(named.named.iter())
                .map(|(ident, f)| gen_visit_field(&ident.to_string(), ident, f))
                .collect::<syn::Result<proc_macro2::TokenStream>>()?;
            Ok(quote_spanned! {input.fields.span()=>
                let Self { #(#names),* } = self;
                #visit_fields
            })
        }
        Fields::Unnamed(unnamed) => {
            let names = fields_unnamed_names(unnamed)?;
            let visit_fields = names
                .iter()
                .enumerate()
                .zip(unnamed.unnamed.iter())
                .map(|((i, ident), field)| gen_visit_field(&i.to_string(), ident, field))
                .collect::<syn::Result<proc_macro2::TokenStream>>()?;
            Ok(quote_spanned! {input.fields.span()=>
                let Self(#(#names),*) = self;
                #visit_fields
            })
        }
        Fields::Unit => Ok(quote_spanned! {input.fields.span()=>}),
    }
}

fn gen_visit_field(
    label: &str,
    ident: &Ident,
    field: &Field,
) -> syn::Result<proc_macro2::TokenStream> {
    let attrs = extract_attrs(&field.attrs)?;
    let field_key = allocative_key(label);
    if attrs.skip {
        Ok(quote_spanned! {field.span()=>})
    } else if let Some(visit) = attrs.visit {
        let ty = &field.ty;
        Ok(quote_spanned! {field.span()=>
            // TODO(nga): figure out how to put this snippet in a member function of the visitor.
            {
                let mut visitor = visitor.enter(#field_key, std::mem::size_of::<#ty>());
                #visit(#ident, &mut visitor);
                visitor.exit();
            }
        })
    } else {
        // Specify type parameter explicitly to prevent implicit conversion.
        let ty = &field.ty;
        Ok(quote_spanned! {ident.span()=>
            visitor.visit_field::<#ty>(#field_key, #ident);
        })
    }
}

#[derive(Default)]
struct AllocativeAttrs {
    skip: bool,
    bound: Option<String>,
    visit: Option<Path>,
}

/// Parse an `#[allocative(...)]` annotation.
#[cfg_attr(feature = "gazebo_lint", allow(gazebo_lint_impl_dupe))] // The custom_keyword macro
fn extract_attrs(attrs: &[Attribute]) -> syn::Result<AllocativeAttrs> {
    syn::custom_keyword!(skip);
    syn::custom_keyword!(bound);
    syn::custom_keyword!(visit);

    let mut opts = AllocativeAttrs::default();

    for attr in attrs.iter() {
        if !attr.path().is_ident("allocative") {
            continue;
        }

        attr.parse_args_with(|input: ParseStream| {
            loop {
                if input.parse::<skip>().is_ok() {
                    if opts.skip {
                        return Err(input.error("`skip` was set twice"));
                    }
                    opts.skip = true;
                } else if input.parse::<bound>().is_ok() {
                    input.parse::<Token![=]>()?;
                    let bound = input.parse::<LitStr>()?;
                    if opts.bound.is_some() {
                        return Err(input.error("`bound` was set twice"));
                    }
                    opts.bound = Some(bound.value());
                } else if input.parse::<visit>().is_ok() {
                    input.parse::<Token![=]>()?;
                    let visit = input.parse::<Path>()?;
                    if opts.visit.is_some() {
                        return Err(input.error("`visit` was set twice"));
                    }
                    opts.visit = Some(visit);
                }

                if input.is_empty() {
                    return Ok(());
                }

                input.parse::<Token![,]>()?;
            }
        })?;
    }

    Ok(opts)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use proc_macro::TokenStream;

mod derive_allocative;
mod root;

#[proc_macro_derive(Allocative, attributes(allocative))]
pub fn derive_allocative(input: TokenStream) -> TokenStream {
    derive_allocative::derive_allocative(input)
}

#[proc_macro_attribute]
pub fn root(attr: TokenStream, input: TokenStream) -> TokenStream {
    root::root(attr, input)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use proc_macro::TokenStream;
use quote::quote_spanned;
use syn::spanned::Spanned;

pub(crate) fn root(_attr: TokenStream, input: TokenStream) -> TokenStream {
    match global_root_impl(input.into()) {
        Ok(tokens) => tokens.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

fn global_root_impl(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2::TokenStream> {
    let item_static = syn::parse2::<syn::ItemStatic>(input)?;

    let name = &item_static.ident;

    let register_name = syn::Ident::new(
        &format!("_{}_register_root", name.to_string().to_lowercase()),
        name.span(),
    );

    Ok(quote_spanned! {item_static.span()=>
        #item_static

        #[allocative::__macro_refs::ctor::ctor]
        fn #register_name() {
            allocative::register_root(&#name);
        }
    })
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![feature(async_closure)]
// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

use std::fs;
use std::io;
use std::path::PathBuf;
use std::sync::Arc;

use buck2::exec;
use buck2::panic;
use buck2::process_context::ProcessContext;
use buck2::TracingLogFile;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::restarter::Restarter;
use buck2_client_ctx::stdin::Stdin;
use buck2_client_ctx::stdio;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::working_dir::WorkingDir;
use buck2_core::logging::init_tracing_for_writer;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use fbinit::FacebookInit;

// fbcode likes to set its own allocator in fbcode.default_allocator
// So when we set our own allocator, buck build buck2 or buck2 build buck2 often breaks.
// Making jemalloc the default only when we do a cargo build.
#[global_allocator]
#[cfg(all(
    any(target_os = "linux", target_os = "macos"),
    not(fbcode_build),
    not(buck_oss_build)
))]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

fn init_logging(_fb: FacebookInit) -> anyhow::Result<Arc<dyn LogConfigurationReloadHandle>> {
    static ENV_TRACING_LOG_FILE_PATH: &str = "BUCK_LOG_TO_FILE_PATH";

    let handle = match std::env::var_os(ENV_TRACING_LOG_FILE_PATH) {
        Some(path) => {
            let path = PathBuf::from(path);
            // we set the writer to stderr first until later, when we have the logdir, set the
            // tracing log sink to that file

            fs::create_dir_all(&path)?;
            let tracing_log = path.join("tracing_log");
            let file = TracingLogFile::new(tracing_log)?;
            init_tracing_for_writer(file)
        }
        _ => init_tracing_for_writer(io::stderr),
    }?;

    #[cfg(fbcode_build)]
    {
        use buck2_client_ctx::subscribers::should_upload_log;
        use buck2_events::sink::scribe;
        use gflags::GflagValue;

        // There are two sources of log spew when building buck2 with Buck and linking against fbcode:
        //   1. folly/logging/xlog, which can be configured via a special configuration string, which we use to
        //      log only critical-level logs. https://github.com/facebook/folly/blob/master/folly/logging/docs/Config.md
        //   2. google log (glog), which is older but still used, which can configured using a flag at runtime.
        //
        // This first call handles the folly config.
        logging::update_logging_config(_fb, "CRITICAL");
        gflags::set_gflag_value(_fb, "minloglevel", GflagValue::U32(5))?;
        gflags::set_gflag_value(_fb, "stderrthreshold", GflagValue::U32(5))?;

        if !should_upload_log()? {
            scribe::disable();
        }
    }

    Ok(handle)
}

// When using a cargo build, some essential services (e.g. RE, scribe)
// fall back to slow paths that give terrible performance.
// Therefore, if we are using cargo, warn strongly.
fn check_cargo() {
    if !cfg!(fbcode_build) && !buck2_core::is_open_source() {
        eprintln!("=====================================================================");
        eprintln!("WARNING: You are using Buck v2 compiled with `cargo`, not `buck`.");
        eprintln!("         Some operations may go slower and logging may be impaired.");
        eprintln!("=====================================================================");
        eprintln!();
    }
}

fn print_retry() -> anyhow::Result<()> {
    buck2_client_ctx::eprintln!("============================================================")?;
    buck2_client_ctx::eprintln!("|| Buck2 has detected that it needs to restart to proceed ||")?;
    buck2_client_ctx::eprintln!("|| Your command will now restart.                         ||")?;
    buck2_client_ctx::eprintln!("============================================================")?;
    buck2_client_ctx::eprintln!()?;
    Ok(())
}

// As this main() is used as the entry point for the `buck daemon` command,
// it must be single-threaded. Commands that want to be multi-threaded/async
// will start up their own tokio runtime.
#[fbinit::main]
fn main(init: fbinit::FacebookInit) -> ! {
    buck2_analysis::init_late_bindings();
    buck2_anon_target::init_late_bindings();
    buck2_action_impl::init_late_bindings();
    buck2_audit_server::init_late_bindings();
    buck2_build_api::init_late_bindings();
    buck2_transition::init_late_bindings();
    buck2_build_signals_impl::init_late_bindings();
    buck2_bxl::init_late_bindings();
    buck2_configured::init_late_bindings();
    buck2_query_impls::init_late_bindings();
    buck2_interpreter_for_build::init_late_bindings();
    buck2_server_commands::init_late_bindings();

    fn main_with_result(init: fbinit::FacebookInit) -> ExitResult {
        panic::initialize(init)?;
        check_cargo();

        static FORCE_WANT_RESTART: EnvHelper<bool> = EnvHelper::new("FORCE_WANT_RESTART");

        let force_want_restart = FORCE_WANT_RESTART.get_copied()?.unwrap_or_default();

        let log_reload_handle = init_logging(init)?;

        let args = std::env::args().collect::<Vec<String>>();
        let cwd = WorkingDir::current_dir()?;
        let mut stdin = Stdin::new()?;
        let mut restarter = Restarter::new();

        let first_trace_id = TraceId::from_env_or_new()?;

        let res = exec(ProcessContext {
            init,
            log_reload_handle: &log_reload_handle,
            stdin: &mut stdin,
            working_dir: &cwd,
            args: &args,
            restarter: &mut restarter,
            trace_id: first_trace_id.dupe(),
            restarted_trace_id: None,
        });

        let restart = |res| {
            if !force_want_restart && !restarter.should_restart() {
                tracing::debug!("No restart was requested");
                return res;
            }

            if stdio::has_written_to_stdout() {
                tracing::debug!("Cannot restart: wrote to stdout");
                return res;
            }

            if print_retry().is_err() {
                tracing::debug!("Cannot restart: warning message cannot be printed");
                return res;
            }

            exec(ProcessContext {
                init,
                log_reload_handle: &log_reload_handle,
                stdin: &mut stdin,
                working_dir: &cwd,
                args: &args,
                restarter: &mut restarter,
                trace_id: TraceId::new(),
                restarted_trace_id: Some(first_trace_id),
            })
        };

        if force_want_restart {
            restart(res)
        } else {
            res.or_else(restart)
        }
    }

    main_with_result(init).report()
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#[cfg(windows)]
pub(crate) fn check_user_allowed() -> anyhow::Result<()> {
    use std::env;
    use std::ffi::OsStr;
    use std::io;
    use std::mem;
    use std::mem::MaybeUninit;
    use std::ptr;

    use anyhow::Context;
    use buck2_wrapper_common::winapi_handle::WinapiHandle;
    use winapi::ctypes::c_void;
    use winapi::shared::minwindef::DWORD;
    use winapi::um::processthreadsapi::GetCurrentProcess;
    use winapi::um::processthreadsapi::OpenProcessToken;
    use winapi::um::securitybaseapi::GetTokenInformation;
    use winapi::um::winnt::TokenElevation;
    use winapi::um::winnt::TOKEN_ELEVATION;
    use winapi::um::winnt::TOKEN_QUERY;

    #[derive(Debug, thiserror::Error)]
    enum CheckUserAllowedError {
        #[error("OpenProcessToken returned null token handle (unreachable)")]
        NullTokenHandle,
    }

    let mut handle = ptr::null_mut();
    let token_ok = unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut handle) };
    if token_ok == 0 {
        return Err(io::Error::last_os_error()).context("OpenProcessToken failed");
    }
    if handle.is_null() {
        return Err(CheckUserAllowedError::NullTokenHandle.into());
    }

    let handle = unsafe { WinapiHandle::new(handle) };
    let size = mem::size_of::<TOKEN_ELEVATION>();
    let elevation: MaybeUninit<TOKEN_ELEVATION> = MaybeUninit::zeroed();
    let mut ret_size = 0;

    let success_get = unsafe {
        GetTokenInformation(
            handle.handle(),
            TokenElevation,
            elevation.as_ptr() as *mut c_void,
            size as DWORD,
            &mut ret_size,
        )
    };
    if success_get == 0 {
        return Err(io::Error::last_os_error()).context("GetTokenInformation failed");
    }

    let elevation_struct: TOKEN_ELEVATION = unsafe { elevation.assume_init() };
    if elevation_struct.TokenIsElevated == 1 {
        // This environment variable is consistently set by CI providers.
        //
        // - GitHub Actions: https://docs.github.com/en/actions/learn-github-actions/variables#default-environment-variables
        // - GitLab CI/CD: https://docs.gitlab.com/ee/ci/variables/predefined_variables.html
        // - CircleCI: https://circleci.com/docs/variables/#built-in-environment-variables
        // - many others
        //
        // In CI, if buck2 got run from an admin shell, we need not worry that a
        // subsequent invocation might come from a non-admin shell. It almost
        // certainly will not.
        let is_ci = env::var_os("CI").as_deref() == Some(OsStr::new("true"));

        if !is_ci {
            tracing::warn!(
                "You're running buck2 from an admin shell. Invocations from non-admin shells will likely fail going forward. To remediate, run `buck2 clean` in this admin shell, then switch to a non-admin shell."
            );
        }
    }
    Ok(())
}

#[cfg(not(windows))]
pub(crate) fn check_user_allowed() -> anyhow::Result<()> {
    use std::os::unix::fs::MetadataExt;

    use anyhow::Context;
    use buck2_core::fs::fs_util;
    use buck2_core::fs::paths::abs_path::AbsPath;
    use buck2_core::soft_error;

    #[derive(Debug, thiserror::Error)]
    #[error("buck2 is not allowed to run as root (unless home dir is owned by root)")]
    struct RootError;

    if nix::unistd::geteuid().is_root() {
        let home_dir = dirs::home_dir().context("home dir not found")?;
        if let Ok(home_dir) = AbsPath::new(&home_dir) {
            let home_dir_metadata = fs_util::metadata(home_dir)?;
            if home_dir_metadata.uid() != 0 {
                soft_error!("root_not_allowed", RootError.into())?;
            }
        }
    }
    Ok(())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! `buck2 audit` command implementation, both client and server.

// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

use std::thread;

use anyhow::Context as _;
use buck2_audit::AuditCommand;
use buck2_client::args::expand_argfiles_with_context;
use buck2_client::commands::build::BuildCommand;
use buck2_client::commands::bxl::BxlCommand;
use buck2_client::commands::clean::CleanCommand;
use buck2_client::commands::ctargets::ConfiguredTargetsCommand;
use buck2_client::commands::debug::DebugCommand;
use buck2_client::commands::init::InitCommand;
use buck2_client::commands::install::InstallCommand;
use buck2_client::commands::kill::KillCommand;
use buck2_client::commands::killall::KillallCommand;
use buck2_client::commands::log::LogCommand;
use buck2_client::commands::lsp::LspCommand;
use buck2_client::commands::profile::ProfileCommand;
use buck2_client::commands::query::aquery::AqueryCommand;
use buck2_client::commands::query::cquery::CqueryCommand;
use buck2_client::commands::query::uquery::UqueryCommand;
use buck2_client::commands::rage::RageCommand;
use buck2_client::commands::root::RootCommand;
use buck2_client::commands::run::RunCommand;
use buck2_client::commands::server::ServerCommand;
use buck2_client::commands::status::StatusCommand;
use buck2_client::commands::subscribe::SubscribeCommand;
use buck2_client::commands::targets::TargetsCommand;
use buck2_client::commands::test::TestCommand;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::cleanup_ctx::AsyncCleanupContextGuard;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::immediate_config::ImmediateConfigContext;
use buck2_client_ctx::streaming::BuckSubcommand;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_client_ctx::tokio_runtime_setup::client_tokio_runtime;
use buck2_client_ctx::version::BuckVersion;
use buck2_common::invocation_paths::InvocationPaths;
use buck2_common::invocation_roots::find_invocation_roots;
use buck2_common::result::ToSharedResultExt;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_event_observer::verbosity::Verbosity;
use buck2_starlark::StarlarkCommand;
use clap::AppSettings;
use clap::Parser;
use dice::DetectCycles;
use dice::WhichDice;
use dupe::Dupe;
use gazebo::variants::VariantName;

use crate::check_user_allowed::check_user_allowed;
use crate::commands::daemon::DaemonCommand;
use crate::commands::docs::DocsCommand;
use crate::commands::forkserver::ForkserverCommand;
use crate::commands::internal_test_runner::InternalTestRunnerCommand;
use crate::process_context::ProcessContext;

#[macro_use]
pub mod panic;
mod check_user_allowed;

pub mod commands;
pub mod process_context;

fn parse_isolation_dir(s: &str) -> anyhow::Result<FileNameBuf> {
    FileNameBuf::try_from(s.to_owned()).context("isolation dir must be a directory name")
}

pub use buck2_server_ctx::logging::TracingLogFile;

/// Options of `buck2` command, before subcommand.
#[derive(Clone, Debug, clap::Parser)]
struct BeforeSubcommandOptions {
    /// Instances of Buck2 share a daemon if and only if their isolation directory is identical.
    /// The isolation directory also influences the output paths provided by Buck2,
    /// and as a result using a non-default isolation dir will cause cache misses (and slower builds).
    #[clap(
        parse(try_from_str = parse_isolation_dir),
        env("BUCK_ISOLATION_DIR"),
        long,
        default_value="v2"
    )]
    isolation_dir: FileNameBuf,

    // TODO: Those should be on the daemon subcommand.
    #[clap(flatten)]
    daemon: DaemonBeforeSubcommandOptions,

    /// How verbose buck should be while logging.
    /// Values:
    /// 0 = Quiet, errors only;
    /// 1 = default;
    /// 2 = more info about errors;
    /// 3 = more info about everything
    #[clap(
        short = 'v',
        long = "verbose",
        default_value = "1",
        global = true,
        value_name = "NUMBER",
        parse(try_from_str = Verbosity::try_from_cli)
    )]
    verbosity: Verbosity,

    /// Do not launch a daemon process, run buck server in client process.
    ///
    /// Note even when running in no-buckd mode, it still writes state files.
    /// In particular, this command effectively kills buckd process
    /// running with the same isolation directory.
    ///
    /// This is an unsupported option used only for development work.
    #[clap(long, global(true), hidden(true))]
    no_buckd: bool,

    /// Print buck wrapper help.
    #[clap(skip)] // @oss-enable
    // @oss-disable: #[clap(long)]
    help_wrapper: bool,
}

#[derive(Clone, Debug, clap::Parser)]
struct DaemonBeforeSubcommandOptions {
    #[clap(env("DICE_DETECT_CYCLES_UNSTABLE"), long, hidden(true))]
    detect_cycles: Option<DetectCycles>,

    #[clap(env("WHICH_DICE_UNSTABLE"), long, hidden(true))]
    which_dice: Option<WhichDice>,

    #[clap(env("ENABLE_TRACE_IO"), long, hidden(true))]
    enable_trace_io: bool,

    /// If passed a given materializer identity, if the materializer state DB matches that
    /// identity, the daemon will not use it and will instead create a new empty materializer
    /// state.
    #[clap(long, hidden(true))]
    reject_materializer_state: Option<String>,
}

#[rustfmt::skip] // Formatting in internal and in OSS versions disagree after oss markers applied.
fn help() -> &'static str {
    concat!(
        "A build system\n",
        "\n",
        "Documentation: https://buck2.build/docs/\n", // @oss-enable
        // @oss-disable: "Documentation: https://internalfb.com/intern/staticdocs/buck2/docs/\n",
    )
}

#[derive(Debug, clap::Parser)]
#[clap(
    name = "buck2",
    about(Some(help())),
    version(BuckVersion::get_version())
)]
pub(crate) struct Opt {
    #[clap(flatten)]
    common_opts: BeforeSubcommandOptions,
    #[clap(subcommand)]
    cmd: CommandKind,
}

impl Opt {
    pub(crate) fn exec(
        self,
        process: ProcessContext<'_>,
        immediate_config: &ImmediateConfigContext,
        matches: &clap::ArgMatches,
        argv: Argv,
    ) -> ExitResult {
        let subcommand_matches = match matches.subcommand().map(|s| s.1) {
            Some(submatches) => submatches,
            None => panic!("Parsed a subcommand but couldn't extract subcommand argument matches"),
        };

        self.cmd.exec(
            process,
            immediate_config,
            subcommand_matches,
            argv,
            self.common_opts,
        )
    }
}

pub fn exec(process: ProcessContext<'_>) -> ExitResult {
    let mut immediate_config = ImmediateConfigContext::new(process.working_dir);
    let mut expanded_args =
        expand_argfiles_with_context(process.args.to_vec(), &mut immediate_config)
            .context("Error expanding argsfiles")?;

    // Override arg0 in `buck2 help`.
    static BUCK2_ARG0: EnvHelper<String> = EnvHelper::new("BUCK2_ARG0");
    if let Some(arg0) = BUCK2_ARG0.get()? {
        expanded_args[0] = arg0.clone();
    }

    let clap = Opt::clap();
    let matches = clap.get_matches_from(&expanded_args);
    let opt: Opt = Opt::from_clap(&matches);

    if opt.common_opts.help_wrapper {
        return ExitResult::err(anyhow::anyhow!(
            "`--help-wrapper` should have been handled by the wrapper"
        ));
    }

    match &opt.cmd {
        CommandKind::Clean(..) | CommandKind::Daemon(..) | CommandKind::Forkserver(..) => {}
        _ => {
            check_user_allowed()?;
        }
    }

    let argv = Argv {
        argv: process.args.to_vec(),
        expanded_argv: expanded_args,
    };

    opt.exec(process, &immediate_config, &matches, argv)
}

#[derive(Debug, clap::Subcommand, VariantName)]
pub(crate) enum CommandKind {
    #[clap(setting(AppSettings::Hidden))]
    Daemon(DaemonCommand),
    #[clap(setting(AppSettings::Hidden))]
    Forkserver(ForkserverCommand),
    #[clap(setting(AppSettings::Hidden))]
    InternalTestRunner(InternalTestRunnerCommand),
    #[clap(subcommand)]
    Audit(AuditCommand),
    Aquery(AqueryCommand),
    Build(BuildCommand),
    Bxl(BxlCommand),
    Test(TestCommand),
    Cquery(CqueryCommand),
    Init(InitCommand),
    Install(InstallCommand),
    Kill(KillCommand),
    Killall(KillallCommand),
    Root(RootCommand),
    /// Alias for `uquery`.
    Query(UqueryCommand),
    Run(RunCommand),
    Server(ServerCommand),
    Status(StatusCommand),
    #[clap(subcommand)]
    Starlark(StarlarkCommand),
    Targets(TargetsCommand),
    Ctargets(ConfiguredTargetsCommand),
    Uquery(UqueryCommand),
    #[clap(subcommand, setting(AppSettings::Hidden))]
    Debug(DebugCommand),
    Docs(DocsCommand),
    #[clap(subcommand)]
    Profile(ProfileCommand),
    #[clap(hide(true))] // @oss-enable
    Rage(RageCommand),
    Clean(CleanCommand),
    #[clap(subcommand)]
    Log(LogCommand),
    Lsp(LspCommand),
    Subscribe(SubscribeCommand),
}

impl CommandKind {
    fn command_name(&self) -> String {
        // clap derive does not expose command name, so do this with gazebo.
        self.variant_name().to_lowercase()
    }

    pub(crate) fn exec(
        self,
        process: ProcessContext<'_>,
        immediate_config: &ImmediateConfigContext,
        matches: &clap::ArgMatches,
        argv: Argv,
        common_opts: BeforeSubcommandOptions,
    ) -> ExitResult {
        let roots = find_invocation_roots(process.working_dir.path());
        let paths = roots
            .map(|r| InvocationPaths {
                roots: r,
                isolation: common_opts.isolation_dir.clone(),
            })
            .shared_error();

        // Handle the daemon command earlier: it wants to fork, but the things we do below might
        // want to create threads.
        if let CommandKind::Daemon(cmd) = self {
            return cmd
                .exec(
                    process.init,
                    process.log_reload_handle.dupe(),
                    paths?,
                    common_opts.daemon,
                    false,
                    || {},
                )
                .into();
        }

        let runtime = client_tokio_runtime()?;
        let async_cleanup = AsyncCleanupContextGuard::new(&runtime);

        let start_in_process_daemon: Option<Box<dyn FnOnce() -> anyhow::Result<()> + Send + Sync>> =
            if common_opts.no_buckd {
                let daemon_startup_config = immediate_config.daemon_startup_config()?.clone();
                let paths = paths.clone()?;
                // Create a function which spawns an in-process daemon.
                Some(Box::new(move || {
                    let (tx, rx) = std::sync::mpsc::channel();
                    // Spawn a thread which runs the daemon.
                    thread::spawn(move || {
                        let tx_clone = tx.clone();
                        let result = DaemonCommand::new_in_process(daemon_startup_config).exec(
                            process.init,
                            <dyn LogConfigurationReloadHandle>::noop(),
                            paths,
                            common_opts.daemon,
                            true,
                            move || drop(tx_clone.send(Ok(()))),
                        );
                        // Since `tx` is unbounded, there's race here: it is possible
                        // that error message will be lost in the channel and not reported anywhere.
                        // Not an issue practically, because daemon does not usually error
                        // after it started listening.
                        if let Err(e) = tx.send(result) {
                            match e.0 {
                                Ok(()) => drop(buck2_client_ctx::eprintln!(
                                    "In-process daemon gracefully stopped"
                                )),
                                Err(e) => drop(buck2_client_ctx::eprintln!(
                                    "In-process daemon run failed: {:#}",
                                    e
                                )),
                            }
                        }
                    });
                    // Wait for listener to start (or to fail).
                    match rx.recv() {
                        Ok(r) => r,
                        Err(_) => Err(anyhow::anyhow!(
                            "In-process daemon failed to start and we don't know why"
                        )),
                    }
                }))
            } else {
                None
            };

        let sanitized_argv = self.sanitize_argv(argv);

        let command_ctx = ClientCommandContext {
            init: process.init,
            immediate_config,
            paths,
            verbosity: common_opts.verbosity,
            start_in_process_daemon,
            command_name: self.command_name(),
            working_dir: process.working_dir.clone(),
            trace_id: process.trace_id.dupe(),
            async_cleanup: async_cleanup.ctx().dupe(),
            stdin: process.stdin,
            restarter: process.restarter,
            restarted_trace_id: process.restarted_trace_id.dupe(),
            sanitized_argv,
            runtime: &runtime,
        };

        match self {
            CommandKind::Daemon(..) => unreachable!("Checked earlier"),
            CommandKind::Forkserver(cmd) => cmd
                .exec(matches, command_ctx, process.log_reload_handle.dupe())
                .into(),
            CommandKind::InternalTestRunner(cmd) => cmd.exec(matches, command_ctx).into(),
            CommandKind::Aquery(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Build(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Bxl(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Test(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Cquery(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Kill(cmd) => cmd.exec(matches, command_ctx).into(),
            CommandKind::Killall(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Clean(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Root(cmd) => cmd.exec(matches, command_ctx).into(),
            CommandKind::Query(cmd) => {
                buck2_client_ctx::eprintln!(
                    "WARNING: \"buck2 query\" is an alias for \"buck2 uquery\". Consider using \"buck2 cquery\" or \"buck2 uquery\" explicitly."
                )?;
                cmd.exec(matches, command_ctx)
            }
            CommandKind::Server(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Status(cmd) => cmd.exec(matches, command_ctx).into(),
            CommandKind::Targets(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Ctargets(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Audit(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Starlark(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Run(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Uquery(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Debug(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Docs(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Profile(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Rage(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Init(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Install(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Log(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Lsp(cmd) => cmd.exec(matches, command_ctx),
            CommandKind::Subscribe(cmd) => cmd.exec(matches, command_ctx),
        }
    }

    fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        match self {
            CommandKind::Daemon(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Forkserver(cmd) => cmd.sanitize_argv(argv),
            CommandKind::InternalTestRunner(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Aquery(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Build(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Bxl(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Test(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Cquery(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Kill(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Killall(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Clean(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Root(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Query(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Server(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Status(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Targets(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Ctargets(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Audit(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Starlark(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Run(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Uquery(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Debug(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Docs(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Profile(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Rage(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Init(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Install(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Log(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Lsp(cmd) => cmd.sanitize_argv(argv),
            CommandKind::Subscribe(cmd) => cmd.sanitize_argv(argv),
        }
    }
}

#[cfg(test)]
mod tests {
    use buck2_client::commands::kill::KillCommand;

    use crate::CommandKind;

    #[test]
    fn test_command_name() {
        assert_eq!("kill", CommandKind::Kill(KillCommand {}).command_name());
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! The panic hook, shared by the buck2 CLI and daemon.
//!
//! The Buck2 CLI unwinds the stack on a panic, while the Buck2 daemon terminates without unwinding.
//! This module sets up a shared panic hook to run before unwinding/termination for crash reporting.

use std::panic;
use std::panic::PanicInfo;

use anyhow::Context as _;
use fbinit::FacebookInit;

/// Initializes the panic hook.
pub fn initialize(fb: FacebookInit) -> anyhow::Result<()> {
    let hook = panic::take_hook();
    panic::set_hook(Box::new(move |info| {
        the_panic_hook(fb, info);
        hook(info);
    }));
    buck2_core::error::initialize(Box::new(move |category, err, loc, options| {
        imp::write_soft_error(
            fb,
            category,
            err,
            buck2_data::Location {
                file: loc.0.to_owned(),
                line: loc.1,
                column: loc.2,
            },
            options,
        );
    }))
    .context("Error initializing soft errors")?;
    Ok(())
}

/// The panic hook, initialized during `initialize`. Invoked immediately on a panic, but prior to the process being
/// aborted.
///
/// The panic hook is called on the same thread that the panic occurred. It is possible to perform a backtrace here
/// to collect additional information.
fn the_panic_hook(fb: FacebookInit, info: &PanicInfo) {
    imp::write_panic_to_scribe(fb, info);
}

mod imp {
    use std::collections::HashMap;
    use std::panic::PanicInfo;
    use std::thread;
    use std::time::Duration;

    use backtrace::Backtrace;
    use buck2_core::error::StructuredErrorOptions;
    use buck2_data::Location;
    use buck2_events::metadata;
    use buck2_events::sink::scribe::new_thrift_scribe_sink_if_enabled;
    use buck2_events::BuckEvent;
    use fbinit::FacebookInit;
    use tokio::runtime::Builder;

    fn get_stack() -> Vec<buck2_data::structured_error::StackFrame> {
        fn ptr_to_string<T>(ptr: *mut T) -> String {
            format!("0x{:x}", ptr as usize)
        }

        let trace = Backtrace::new();
        trace
            .frames()
            .iter()
            .map(|frame| {
                let symbols = frame
                    .symbols()
                    .iter()
                    .map(|symbol| buck2_data::structured_error::Symbol {
                        name: symbol
                            .name()
                            .map_or_else(|| "".to_owned(), |s| s.to_string()),
                        address: symbol.addr().map_or_else(|| "".to_owned(), ptr_to_string),
                        file: symbol
                            .filename()
                            .map_or_else(|| "".to_owned(), |p| p.display().to_string()),
                        line: symbol.lineno().unwrap_or(0),
                        column: symbol.colno().unwrap_or(0),
                    })
                    .collect::<Vec<_>>();

                buck2_data::structured_error::StackFrame {
                    instruction_pointer: ptr_to_string(frame.ip()),
                    symbol_address: ptr_to_string(frame.symbol_address()),
                    module_base_address: frame
                        .module_base_address()
                        .map_or_else(|| "".to_owned(), ptr_to_string),
                    symbols,
                }
            })
            .collect()
    }

    /// Extracts a stringly-formatted payload from the given PanicInfo - usually the argument to `panic!`.
    fn get_message_for_panic(info: &PanicInfo) -> String {
        // Rust panic payloads can be anything, but they generally take one of two forms:
        //  1. &'static str, for panics with a constant string parameter,
        //  2. String, for panics that use the format `{}` syntax to construct a message.
        //
        // Since PanicInfo's Display implementation is implemented in libcore, it can't cover the String case (which is)
        // a liballoc exclusive), so this code here checks for formatted messages and uses that as the message if present.
        if let Some(literal_msg) = info.payload().downcast_ref::<&str>() {
            (*literal_msg).to_owned()
        } else if let Some(format_msg) = info.payload().downcast_ref::<String>() {
            format_msg.clone()
        } else {
            "explicit panic with no message".to_owned()
        }
    }

    /// Collects metadata from the current environment for use in LogView.
    fn get_metadata_for_panic() -> HashMap<String, String> {
        let mut map = metadata::collect();
        if let Some(commands) = buck2_server::active_commands::try_active_commands() {
            let commands = commands.keys().map(|id| id.to_string()).collect::<Vec<_>>();
            map.insert("active_commands".to_owned(), commands.join(","));
        }
        map
    }

    /// Writes a representation of the given `PanicInfo` to Scribe, via the `Panic` event.
    pub(crate) fn write_panic_to_scribe(fb: FacebookInit, info: &PanicInfo) {
        let message = get_message_for_panic(info);
        let location = info.location().map(|loc| Location {
            file: loc.file().to_owned(),
            line: loc.line(),
            column: loc.column(),
        });
        write_to_scribe(
            fb,
            panic_payload(location, message, get_stack(), &Default::default(), None),
        );
    }

    pub(crate) fn write_soft_error(
        fb: FacebookInit,
        category: &str,
        err: &anyhow::Error,
        location: Location,
        options: StructuredErrorOptions,
    ) {
        let event = panic_payload(
            Some(location),
            format!("Soft Error: {}: {:#}", category, err),
            Vec::new(),
            &options,
            Some(category),
        );

        // If the soft error was fired in a context with an ambient dispatcher, then we only send
        // it there, but some contexts don't have one, and in that case, we notify all running
        // commands.
        match buck2_events::dispatch::get_dispatcher_opt() {
            Some(dispatcher) => {
                dispatcher.instant_event(event.clone());
            }
            None => {
                if !buck2_server::active_commands::broadcast_instant_event(&event) && !options.quiet
                {
                    tracing::warn!("Warning \"{}\": {:#}", category, err);
                }
            }
        }

        write_to_scribe(fb, event);
    }

    fn panic_payload(
        location: Option<Location>,
        message: String,
        backtrace: Vec<buck2_data::structured_error::StackFrame>,
        options: &StructuredErrorOptions,
        soft_error_category: Option<&str>,
    ) -> buck2_data::StructuredError {
        let metadata = get_metadata_for_panic();
        buck2_data::StructuredError {
            location,
            payload: message,
            metadata,
            backtrace,
            quiet: options.quiet,
            task: Some(options.task),
            soft_error_category: soft_error_category.map(ToOwned::to_owned),
            daemon_in_memory_state_is_corrupted: options.daemon_in_memory_state_is_corrupted,
            daemon_materializer_state_is_corrupted: options.daemon_materializer_state_is_corrupted,
            action_cache_is_corrupted: options.action_cache_is_corrupted,
        }
    }

    /// Writes a representation of the given error (hard or soft) to Scribe
    fn write_to_scribe(fb: FacebookInit, data: buck2_data::StructuredError) {
        use std::time::SystemTime;

        use buck2_core::facebook_only;
        use buck2_data::InstantEvent;
        use buck2_events::sink::scribe;
        use buck2_wrapper_common::invocation_id::TraceId;

        facebook_only();
        if !scribe::is_enabled() {
            return;
        }

        let sink = match new_thrift_scribe_sink_if_enabled(
            fb,
            /* buffer size */ 100,
            /* retry_backoff */ Duration::from_millis(500),
            /* retry_attempts */ 5,
            /* message_batch_size */ None,
        ) {
            Ok(Some(sink)) => sink,
            _ => {
                // We're already panicking and we can't connect to the scribe daemon? Things are bad and we're SOL.
                return;
            }
        };

        // There are some dubious ways of acquiring a handle to the current Tokio runtime, if one exists, such as this
        // one: https://docs.rs/tokio/latest/tokio/runtime/struct.Handle.html#method.current. However, there doesn't
        // appear to be a good way to figure out if the thread we're running on is a tokio thread at all.
        //
        // Since this is the panic handler, and the panic handler runs in the context of the thread that panicked, we
        // can make no assumptions about whether the thread we're on is a tokio thread. To flush the Scribe client,
        // which is normally an async operation, we spawn a new thread, spawn a runtime on that, and await the flush
        // on that thread.
        //
        // Note that if we fail to spawn a writer thread, then we just won't log.
        let _err = thread::Builder::new()
            .spawn(move || {
                let runtime = Builder::new_current_thread().enable_all().build().unwrap();
                runtime.block_on(
                    sink.send_now(BuckEvent::new(
                        SystemTime::now(),
                        TraceId::new(),
                        None,
                        None,
                        InstantEvent {
                            data: Some(data.into()),
                        }
                        .into(),
                    )),
                );
            })
            .map_err(|_| ())
            .and_then(|t| t.join().map_err(|_| ()));
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::sync::Arc;

use buck2_client_ctx::restarter::Restarter;
use buck2_client_ctx::stdin::Stdin;
use buck2_core::fs::working_dir::WorkingDir;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_wrapper_common::invocation_id::TraceId;

/// State passed down from `main` to this crate.
pub struct ProcessContext<'a> {
    pub init: fbinit::FacebookInit,
    pub log_reload_handle: &'a Arc<dyn LogConfigurationReloadHandle>,
    pub stdin: &'a mut Stdin,
    pub working_dir: &'a WorkingDir,
    pub args: &'a [String],
    pub restarter: &'a mut Restarter,
    pub trace_id: TraceId,
    /// An invocation that this invocation is a restart of.
    pub restarted_trace_id: Option<TraceId>,
}