// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build !tinygo
// +build !tinygo

package memory

import (
	
	
	
	
	
	
	
	
)

type CheckedAllocator struct {
	mem Allocator
	sz  atomic.Int64

	allocs sync.Map
}

func ( Allocator) *CheckedAllocator {
	return &CheckedAllocator{mem: }
}

func ( *CheckedAllocator) () int { return int(.sz.Load()) }

func ( *CheckedAllocator) ( int) []byte {
	.sz.Add(int64())
	 := .mem.Allocate()
	if  == 0 {
		return 
	}

	 := uintptr(unsafe.Pointer(&[0]))
	 := make([]uintptr, maxRetainedFrames)

	// For historical reasons the meaning of the skip argument
	// differs between Caller and Callers. For Callers, 0 identifies
	// the frame for the caller itself. We skip 2 additional frames
	// here to get to the caller right before the call to Allocate.
	runtime.Callers(allocFrames+2, )
	 := runtime.CallersFrames()
	if , , ,  := runtime.Caller(allocFrames);  {
		.allocs.Store(, &dalloc{pc: , line: , sz: , callersFrames: })
	}
	return 
}

func ( *CheckedAllocator) ( int,  []byte) []byte {
	.sz.Add(int64( - len()))

	 := uintptr(unsafe.Pointer(&[0]))
	 := .mem.Reallocate(, )
	if  == 0 {
		return 
	}

	 := uintptr(unsafe.Pointer(&[0]))
	.allocs.Delete()
	 := make([]uintptr, maxRetainedFrames)

	// For historical reasons the meaning of the skip argument
	// differs between Caller and Callers. For Callers, 0 identifies
	// the frame for the caller itself. We skip 2 additional frames
	// here to get to the caller right before the call to Reallocate.
	runtime.Callers(reallocFrames+2, )
	 := runtime.CallersFrames()
	if , , ,  := runtime.Caller(reallocFrames);  {
		.allocs.Store(, &dalloc{pc: , line: , sz: , callersFrames: })
	}

	return 
}

func ( *CheckedAllocator) ( []byte) {
	.sz.Add(int64(len() * -1))
	defer .mem.Free()

	if len() == 0 {
		return
	}

	 := uintptr(unsafe.Pointer(&[0]))
	.allocs.Delete()
}

// typically the allocations are happening in memory.Buffer, not by consumers calling
// allocate/reallocate directly. As a result, we want to skip the caller frames
// of the inner workings of Buffer in order to find the caller that actually triggered
// the allocation via a call to Resize/Reserve/etc.
const (
	defAllocFrames       = 4
	defReallocFrames     = 3
	defMaxRetainedFrames = 0
)

// Use the environment variables ARROW_CHECKED_ALLOC_FRAMES and ARROW_CHECKED_REALLOC_FRAMES
// to control how many frames it skips when storing the caller for allocations/reallocs
// when using this to find memory leaks. Use ARROW_CHECKED_MAX_RETAINED_FRAMES to control how
// many frames are retained for printing the stack trace of a leak.
var allocFrames, reallocFrames, maxRetainedFrames int = defAllocFrames, defReallocFrames, defMaxRetainedFrames

func init() {
	if ,  := os.LookupEnv("ARROW_CHECKED_ALLOC_FRAMES");  {
		if ,  := strconv.Atoi();  == nil {
			allocFrames = 
		}
	}

	if ,  := os.LookupEnv("ARROW_CHECKED_REALLOC_FRAMES");  {
		if ,  := strconv.Atoi();  == nil {
			reallocFrames = 
		}
	}

	if ,  := os.LookupEnv("ARROW_CHECKED_MAX_RETAINED_FRAMES");  {
		if ,  := strconv.Atoi();  == nil {
			maxRetainedFrames = 
		}
	}
}

type dalloc struct {
	pc            uintptr
	line          int
	sz            int
	callersFrames *runtime.Frames
}

type TestingT interface {
	Errorf(format string, args ...interface{})
	Helper()
}

func ( *CheckedAllocator) ( TestingT,  int) {
	.allocs.Range(func(,  interface{}) bool {
		 := .(*dalloc)
		 := runtime.FuncForPC(.pc)
		 := .callersFrames
		var  strings.Builder
		for {
			,  := .Next()
			if .Line == 0 {
				break
			}
			.WriteString("\t")
			// frame.Func is a useful source of information if it's present.
			// It may be nil for non-Go code or fully inlined functions.
			if  := .Func;  != nil {
				// format as func name + the offset in bytes from func entrypoint
				.WriteString(fmt.Sprintf("%s+%x", .Name(), .PC-.Entry()))
			} else {
				// fallback to outer func name + file line
				.WriteString(fmt.Sprintf("%s, line %d", .Function, .Line))
			}

			// Write a proper file name + line, so it's really easy to find the leak
			.WriteString("\n\t\t")
			.WriteString(.File + ":" + strconv.Itoa(.Line))
			.WriteString("\n")
			if ! {
				break
			}
		}

		,  := .FileLine(.pc)
		.Errorf("LEAK of %d bytes FROM\n\t%s+%x\n\t\t%s:%d\n%v",
			.sz,
			.Name(), .pc-.Entry(), // func name + offset in bytes between frame & entrypoint to func
			, , // a proper file name + line, so it's really easy to find the leak
			.String(),
		)
		return true
	})

	if int(.sz.Load()) !=  {
		.Helper()
		.Errorf("invalid memory size exp=%d, got=%d", , .sz.Load())
	}
}

type CheckedAllocatorScope struct {
	alloc *CheckedAllocator
	sz    int
}

func ( *CheckedAllocator) *CheckedAllocatorScope {
	 := .sz.Load()
	return &CheckedAllocatorScope{alloc: , sz: int()}
}

func ( *CheckedAllocatorScope) ( TestingT) {
	 := int(.alloc.sz.Load())
	if .sz !=  {
		.Helper()
		.Errorf("invalid memory size exp=%d, got=%d", .sz, )
	}
}

var _ Allocator = (*CheckedAllocator)(nil)